diff --git a/sam3.cpp b/sam3.cpp index 43d88a8..c7148a6 100644 --- a/sam3.cpp +++ b/sam3.cpp @@ -18,15 +18,12 @@ #define STB_IMAGE_WRITE_IMPLEMENTATION #ifdef _WIN32 -#include #include #define popen _popen #define pclose _pclose -#define mkdir(path, mode) _mkdir(path) #define SAM3_NULL_DEV "NUL" #define SAM3_POPEN_READ "rb" #else -#include #define SAM3_NULL_DEV "/dev/null" #define SAM3_POPEN_READ "r" #endif @@ -1009,9 +1006,6 @@ static struct ggml_tensor* sam3_layer_norm_2d(struct ggml_context* ctx, struct ggml_tensor* w, struct ggml_tensor* b); -static bool sam3_copy_tensor_to_f32(struct ggml_tensor * t, - std::vector & output); - /***************************************************************************** ** Internal Helper Implementations *****************************************************************************/ @@ -6215,6449 +6209,3311 @@ bool sam3_encode_image(sam3_state& state, return true; } -static void sam3_clear_encoder_state(sam3_state & state) { - state.vit_output = nullptr; - for (int i = 0; i < 4; ++i) { - state.neck_det[i] = nullptr; - state.neck_trk[i] = nullptr; - state.neck_det_pe[i] = nullptr; - state.neck_trk_pe[i] = nullptr; - } -} - -static bool sam3_mark_named_outputs(struct ggml_context * ctx, - const std::vector & output_tensors) { - for (const auto & name : output_tensors) { - struct ggml_tensor * t = ggml_get_tensor(ctx, name.c_str()); - if (!t) { - fprintf(stderr, "%s: requested tensor '%s' was not found\n", __func__, name.c_str()); - return false; - } - ggml_set_output(t); - } - return true; -} - -bool sam3_encode_vit_from_preprocessed_selective(sam3_state & state, - const sam3_model & model, - const float * chw_data, - int img_size, - const std::vector & output_tensors) { - const auto & hp = model.hparams; - - if (img_size != hp.img_size) { - fprintf(stderr, "%s: img_size mismatch: got %d, expected %d\n", - __func__, img_size, hp.img_size); - return false; - } - - state.orig_width = img_size; - state.orig_height = img_size; - - const size_t buf_size = ggml_tensor_overhead() * 4096 + ggml_graph_overhead() * 2; - struct ggml_init_params gparams = { - /*.mem_size =*/ buf_size, - /*.mem_buffer =*/ nullptr, - /*.no_alloc =*/ true, - }; - struct ggml_context * ctx0 = ggml_init(gparams); - if (!ctx0) { - fprintf(stderr, "%s: failed to init compute context\n", __func__); - return false; - } +/***************************************************************************** +** Multi-head attention helper (used by fusion encoder, DETR decoder, seg head) +*****************************************************************************/ - auto * inp = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, img_size, img_size, 3, 1); - ggml_set_name(inp, "input_image"); - ggml_set_input(inp); +// Standard multi-head attention with fused in_proj. +// q_in, k_in, v_in: [D, N, B] (if fused_qkv, only q_in is used and contains QKV stacked) +// in_proj_w: [D, 3*D] (fused Q/K/V projection) +// in_proj_b: [3*D] +// out_proj_w: [D, D], out_proj_b: [D] +// n_heads: number of attention heads +// Returns: [D, N_q, B] +// +// If separate_kv is true, q_in/k_in/v_in are already separate (no fused proj needed). +// The in_proj is applied to form Q from q_in, and K/V from the concatenated k/v source. +static struct ggml_tensor* sam3_multihead_attn_fused( + struct ggml_context* ctx, + struct ggml_tensor* q_in, // [D, N_q, B] + struct ggml_tensor* kv_in, // [D, N_kv, B] (can be same as q_in for self-attn) + struct ggml_tensor* in_proj_w, // [D, 3*D] — fused QKV weights + struct ggml_tensor* in_proj_b, // [3*D] + struct ggml_tensor* out_proj_w, // [D, D] + struct ggml_tensor* out_proj_b, // [D] + int n_heads, + struct ggml_tensor* attn_mask = nullptr) // [N_kv, N_q] or nullptr +{ + const int64_t D = q_in->ne[0]; // 256 + const int64_t N_q = q_in->ne[1]; + const int64_t B = q_in->ne[2]; + const int64_t N_kv = kv_in->ne[1]; + const int64_t HD = D / n_heads; - auto * vit_out = sam3_build_vit_graph(ctx0, inp, model); - ggml_set_name(vit_out, "vit_output"); - ggml_set_output(vit_out); + auto* q_w = ggml_view_2d(ctx, in_proj_w, D, D, in_proj_w->nb[1], 0); + auto* k_w = ggml_view_2d(ctx, in_proj_w, D, D, in_proj_w->nb[1], D * in_proj_w->nb[1]); + auto* v_w = ggml_view_2d(ctx, in_proj_w, D, D, in_proj_w->nb[1], 2 * D * in_proj_w->nb[1]); - if (!sam3_mark_named_outputs(ctx0, output_tensors)) { - ggml_free(ctx0); - return false; - } + auto* q_b = ggml_view_1d(ctx, in_proj_b, D, 0); + auto* k_b = ggml_view_1d(ctx, in_proj_b, D, D * sizeof(float)); + auto* v_b = ggml_view_1d(ctx, in_proj_b, D, 2 * D * sizeof(float)); - struct ggml_cgraph * graph = ggml_new_graph_custom(ctx0, 8192, false); - ggml_build_forward_expand(graph, vit_out); + auto* Q = ggml_add(ctx, ggml_mul_mat(ctx, q_w, q_in), q_b); + auto* K = ggml_add(ctx, ggml_mul_mat(ctx, k_w, kv_in), k_b); + auto* V = ggml_add(ctx, ggml_mul_mat(ctx, v_w, kv_in), v_b); - auto * galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); - if (!ggml_gallocr_reserve(galloc, graph)) { - fprintf(stderr, "%s: failed to reserve graph memory\n", __func__); - ggml_gallocr_free(galloc); - ggml_free(ctx0); - return false; - } + Q = ggml_reshape_4d(ctx, Q, HD, n_heads, N_q, B); + Q = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 1, 3)); // [HD, N_q, NH, B] - if (!ggml_gallocr_alloc_graph(galloc, graph)) { - fprintf(stderr, "%s: failed to allocate graph\n", __func__); - ggml_gallocr_free(galloc); - ggml_free(ctx0); - return false; - } + K = ggml_reshape_4d(ctx, K, HD, n_heads, N_kv, B); + K = ggml_cont(ctx, ggml_permute(ctx, K, 0, 2, 1, 3)); // [HD, N_kv, NH, B] - const size_t data_bytes = (size_t) 3 * img_size * img_size * sizeof(float); - ggml_backend_tensor_set(inp, chw_data, 0, data_bytes); - if (!sam3_graph_compute(model.backend, graph, state.n_threads)) { - ggml_gallocr_free(galloc); - ggml_free(ctx0); - return false; - } + V = ggml_reshape_4d(ctx, V, HD, n_heads, N_kv, B); + V = ggml_permute(ctx, V, 0, 2, 1, 3); // [HD, N_kv, NH, B] non-contiguous; flash_attn uses strides - if (state.galloc) { - ggml_gallocr_free(state.galloc); - } - if (state.ctx) { - ggml_free(state.ctx); - } - if (state.pe_buf) { - ggml_backend_buffer_free(state.pe_buf); - state.pe_buf = nullptr; - } - if (state.pe_ctx) { - ggml_free(state.pe_ctx); - state.pe_ctx = nullptr; - } + float scale = 1.0f / sqrtf((float)HD); + auto* attn_out = ggml_flash_attn_ext(ctx, Q, K, V, attn_mask, scale, 0.0f, 0.0f); - state.ctx = ctx0; - state.galloc = galloc; - state.backend = model.backend; - sam3_clear_encoder_state(state); - state.vit_output = vit_out; - state.pe_cache_valid = false; + auto* merged = ggml_reshape_3d(ctx, attn_out, D, N_q, B); + merged = ggml_mul_mat(ctx, out_proj_w, merged); + merged = ggml_add(ctx, merged, out_proj_b); - return true; + return merged; } -static void sam3_normalize_ne4(const int64_t input_ne[4], int64_t ne[4]); -static struct ggml_tensor * sam3_new_f32_tensor_4d_from_ne(struct ggml_context * ctx, - const int64_t ne[4]); -static bool sam3_copy_tensor_to_f32(struct ggml_tensor * t, std::vector & output); -static struct ggml_tensor * sam3_build_vit_prefix_stage_from_input(struct ggml_context * ctx, - struct ggml_tensor * input, - const sam3_model & model, - sam3_vit_prefix_stage stage); - -bool sam3_test_run_vit_block0_input(const sam3_model & model, - const float * chw_data, - int img_size, - std::vector & output_data, - int64_t output_ne[4], - int n_threads) { - if (!chw_data) { - fprintf(stderr, "%s: chw_data is null\n", __func__); - return false; - } - if (img_size != model.hparams.img_size) { - fprintf(stderr, "%s: img_size mismatch: got %d, expected %d\n", - __func__, img_size, model.hparams.img_size); - return false; +static struct ggml_tensor* sam3_expand_token_attn_bias( + struct ggml_context* ctx, + struct ggml_tensor* token_bias, // [T, 1, B] or nullptr + int64_t n_q, + int n_heads, + int64_t batch) { + if (!token_bias) { + return nullptr; } - const size_t ctx_size = ggml_tensor_overhead() * 256 + ggml_graph_overhead(); - ggml_init_params params = { - /*.mem_size =*/ ctx_size, - /*.mem_buffer =*/ nullptr, - /*.no_alloc =*/ true, - }; - - struct ggml_context * ctx = ggml_init(params); - if (!ctx) { - fprintf(stderr, "%s: failed to create ggml context\n", __func__); - return false; - } + const int64_t n_kv = token_bias->ne[0]; + auto* bias_4d = ggml_reshape_4d(ctx, token_bias, n_kv, 1, 1, batch); + auto* full_bias = ggml_repeat( + ctx, + bias_4d, + ggml_new_tensor_4d(ctx, GGML_TYPE_F32, n_kv, n_q, n_heads, batch)); - struct ggml_tensor * input = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, img_size, img_size, 3, 1); - ggml_set_name(input, "vit_prefix_input"); - ggml_set_input(input); + // ggml_flash_attn_ext expects mask storage in fp16. + return ggml_cont(ctx, ggml_cast(ctx, full_bias, GGML_TYPE_F16)); +} - struct ggml_tensor * output = sam3_build_vit_prefix_graph(ctx, input, model); - ggml_set_name(output, "vit_block0_input"); - ggml_set_output(output); +/***************************************************************************** +** Geometry / exemplar encoder — graph building +*****************************************************************************/ - struct ggml_cgraph * graph = ggml_new_graph_custom(ctx, 1024, false); - ggml_build_forward_expand(graph, output); +// Sinusoidal positional encoding for box coordinates. +// Matches Python PositionEmbeddingSine.encode_boxes(cx, cy, w, h). +// Output: [258] = [pos_y(128), pos_x(128), h, w] +static void sam3_sine_encode_box(float* out, float cx, float cy, float w, float h, + int num_pos_feats, int temperature) { + const float scale = 2.0f * (float)M_PI; + float x_embed = cx * scale; + float y_embed = cy * scale; - auto * galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); - if (!ggml_gallocr_reserve(galloc, graph) || !ggml_gallocr_alloc_graph(galloc, graph)) { - fprintf(stderr, "%s: failed to allocate graph\n", __func__); - ggml_gallocr_free(galloc); - ggml_free(ctx); - return false; + // dim_t[i] = temperature^(2*(i//2)/num_pos_feats) + for (int i = 0; i < num_pos_feats; ++i) { + int div_idx = 2 * (i / 2); + float dim_t = powf((float)temperature, (float)div_idx / (float)num_pos_feats); + float px = x_embed / dim_t; + float py = y_embed / dim_t; + // Interleaved sin/cos: even indices get sin, odd get cos + if (i % 2 == 0) { + out[i] = sinf(py); // pos_y first + out[num_pos_feats + i] = sinf(px); // pos_x second + } else { + out[i] = cosf(py); + out[num_pos_feats + i] = cosf(px); + } } + out[2 * num_pos_feats] = h; // h + out[2 * num_pos_feats + 1] = w; // w +} - const size_t input_bytes = (size_t) 3 * img_size * img_size * sizeof(float); - ggml_backend_tensor_set(input, chw_data, 0, input_bytes); - sam3_graph_compute(model.backend, graph, n_threads); +// CPU-side ROI Align matching torchvision.ops.roi_align behavior. +// Features in ggml [C, W, H] layout. Box in XYXY format scaled to feature grid coords. +// Uses sub-sampling matching torchvision's sampling_ratio=0 (auto). +// Output: [C * roi_size * roi_size] in ggml layout. +static void sam3_roi_align_single( + const float* feats, // [C, W, H] ggml layout (C innermost, then W, then H) + int C, int W_feat, int H_feat, + float x0, float y0, float x1, float y1, + int roi_size, + float* out) // [C, roi_size, roi_size] +{ + float roi_w = std::max(x1 - x0, 1e-6f); + float roi_h = std::max(y1 - y0, 1e-6f); + float bin_w = roi_w / (float)roi_size; + float bin_h = roi_h / (float)roi_size; - for (int i = 0; i < 4; ++i) { - output_ne[i] = output->ne[i]; - } + // Match torchvision sampling_ratio=0: use ceil(bin_size) sub-samples + int sample_y_count = std::max(1, (int)ceilf(bin_h)); + int sample_x_count = std::max(1, (int)ceilf(bin_w)); + float inv_count = 1.0f / (float)(sample_y_count * sample_x_count); - const bool ok = sam3_copy_tensor_to_f32(output, output_data); - ggml_gallocr_free(galloc); - ggml_free(ctx); - return ok; -} - -bool sam3_test_run_vit_prefix_stage(const sam3_model & model, - sam3_vit_prefix_stage stage, - const float * input_data, - const int64_t input_ne[4], - std::vector & output_data, - int64_t output_ne[4], - int n_threads) { - if (!input_data) { - fprintf(stderr, "%s: input_data is null\n", __func__); - return false; - } + auto bilinear_sample = [&](float sx, float sy, int c) -> float { + // Clamp to [-0.5, W-0.5] then adjust — matching torchvision + if (sx < -1.0f || sx > (float)W_feat || sy < -1.0f || sy > (float)H_feat) + return 0.0f; - int64_t ne[4]; - sam3_normalize_ne4(input_ne, ne); + sy = std::max(0.0f, sy); + sx = std::max(0.0f, sx); - switch (stage) { - case SAM3_VIT_PREFIX_STAGE_PATCH_IM2COL: - case SAM3_VIT_PREFIX_STAGE_PATCH_EMBED: - if (ne[0] != model.hparams.img_size || ne[1] != model.hparams.img_size || ne[2] != 3) { - fprintf(stderr, "%s: patch_embed input shape mismatch [%lld,%lld,%lld,%lld]\n", - __func__, - (long long) ne[0], (long long) ne[1], (long long) ne[2], (long long) ne[3]); - return false; - } - break; + int y_lo = (int)sy; + int x_lo = (int)sx; + int y_hi = std::min(y_lo + 1, H_feat - 1); + int x_hi = std::min(x_lo + 1, W_feat - 1); + y_lo = std::min(y_lo, H_feat - 1); + x_lo = std::min(x_lo, W_feat - 1); - case SAM3_VIT_PREFIX_STAGE_PATCH_MULMAT_RAW: - case SAM3_VIT_PREFIX_STAGE_PATCH_MULMAT: - if (ne[0] != model.vit.patch_embed_w->ne[0] * model.vit.patch_embed_w->ne[1] * model.vit.patch_embed_w->ne[2] || - ne[1] != model.hparams.n_img_embd() || - ne[2] != model.hparams.n_img_embd()) { - fprintf(stderr, "%s: patch_mulmat input shape mismatch [%lld,%lld,%lld,%lld]\n", - __func__, - (long long) ne[0], (long long) ne[1], (long long) ne[2], (long long) ne[3]); - return false; - } - break; + float ly = sy - (float)y_lo; + float lx = sx - (float)x_lo; - default: - if (ne[0] != model.hparams.vit_embed_dim || - ne[1] != model.hparams.n_img_embd() || - ne[2] != model.hparams.n_img_embd()) { - fprintf(stderr, "%s: prefix feature input shape mismatch [%lld,%lld,%lld,%lld]\n", - __func__, - (long long) ne[0], (long long) ne[1], (long long) ne[2], (long long) ne[3]); - return false; - } - break; - } + // ggml layout: feats[c + x * C + y * C * W_feat] + float v00 = feats[c + x_lo * C + y_lo * C * W_feat]; + float v10 = feats[c + x_hi * C + y_lo * C * W_feat]; + float v01 = feats[c + x_lo * C + y_hi * C * W_feat]; + float v11 = feats[c + x_hi * C + y_hi * C * W_feat]; - const size_t ctx_size = ggml_tensor_overhead() * 128 + ggml_graph_overhead(); - ggml_init_params params = { - /*.mem_size =*/ ctx_size, - /*.mem_buffer =*/ nullptr, - /*.no_alloc =*/ true, + return v00 * (1 - ly) * (1 - lx) + v10 * (1 - ly) * lx + v01 * ly * (1 - lx) + v11 * ly * lx; }; - struct ggml_context * ctx = ggml_init(params); - if (!ctx) { - fprintf(stderr, "%s: failed to create ggml context\n", __func__); - return false; + for (int ph = 0; ph < roi_size; ++ph) { + for (int pw = 0; pw < roi_size; ++pw) { + for (int c = 0; c < C; ++c) { + float sum = 0.0f; + for (int iy = 0; iy < sample_y_count; ++iy) { + float sy = y0 + bin_h * ((float)ph + ((float)iy + 0.5f) / (float)sample_y_count); + for (int ix = 0; ix < sample_x_count; ++ix) { + float sx = x0 + bin_w * ((float)pw + ((float)ix + 0.5f) / (float)sample_x_count); + sum += bilinear_sample(sx, sy, c); + } + } + out[c + pw * C + ph * C * roi_size] = sum * inv_count; + } + } } +} - struct ggml_tensor * input = nullptr; - if (stage == SAM3_VIT_PREFIX_STAGE_PATCH_MULMAT_RAW || stage == SAM3_VIT_PREFIX_STAGE_PATCH_MULMAT) { - input = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, ne[0], ne[1], ne[2], ne[3]); - } else { - input = sam3_new_f32_tensor_4d_from_ne(ctx, ne); - } - ggml_set_name(input, "vit_prefix_stage_input"); - ggml_set_input(input); +// Build geometry encoder graph and pre-compute box embeddings on CPU. +// Returns the geometry features as a pre-computed input tensor [D, N_geo, 1] +// where N_geo = n_exemplar_boxes + 1 (CLS). +// For dummy prompts (no boxes), N_geo = 1 (just CLS token). +struct sam3_geom_result { + struct ggml_tensor* geo_feats; // [D, N_geo, 1] + int n_tokens; // N_geo +}; - struct ggml_tensor * output = sam3_build_vit_prefix_stage_from_input(ctx, input, model, stage); - if (!output) { - fprintf(stderr, "%s: failed to build prefix stage %d\n", __func__, (int) stage); - ggml_free(ctx); - return false; - } - ggml_set_name(output, "vit_prefix_stage_output"); - ggml_set_output(output); - - struct ggml_cgraph * graph = ggml_new_graph_custom(ctx, 512, false); - ggml_build_forward_expand(graph, output); - - auto * galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); - if (!ggml_gallocr_reserve(galloc, graph) || !ggml_gallocr_alloc_graph(galloc, graph)) { - fprintf(stderr, "%s: failed to allocate graph\n", __func__); - ggml_gallocr_free(galloc); - ggml_free(ctx); - return false; - } - - if (input->type == GGML_TYPE_F16) { - const int64_t numel = ggml_nelements(input); - std::vector tmp((size_t) numel); - ggml_fp32_to_fp16_row(input_data, tmp.data(), numel); - ggml_backend_tensor_set(input, tmp.data(), 0, (size_t) numel * sizeof(ggml_fp16_t)); - } else { - ggml_backend_tensor_set(input, input_data, 0, (size_t) ggml_nbytes(input)); - } - sam3_graph_compute(model.backend, graph, n_threads); - - for (int i = 0; i < 4; ++i) { - output_ne[i] = output->ne[i]; - } - - const bool ok = sam3_copy_tensor_to_f32(output, output_data); - ggml_gallocr_free(galloc); - ggml_free(ctx); - return ok; -} - -bool sam3_test_run_patch_mulmat_host_ref(const sam3_model & model, - const float * input_data, - const int64_t input_ne[4], - bool use_double_accum, - std::vector & output_data, - int64_t output_ne[4]) { - if (!input_data) { - return false; - } - - int64_t ne[4]; - sam3_normalize_ne4(input_ne, ne); +static sam3_geom_result sam3_build_geom_enc_graph( + struct ggml_context* ctx, + const sam3_model& model, + const sam3_pcs_params& params, + struct ggml_tensor* img_feats, // [D, N_img, 1] where N_img = H*H + struct ggml_tensor* img_pe) // [D, N_img, 1] sinusoidal PE +{ + const auto& ge = model.geom_enc; + const int D = model.hparams.neck_dim; // 256 + const int n_heads = 8; + const int n_boxes = (int)(params.pos_exemplars.size() + params.neg_exemplars.size()); + const int N_geo = n_boxes + 1; // +1 for CLS - const int64_t patch_k = model.vit.patch_embed_w->ne[0] * model.vit.patch_embed_w->ne[1] * model.vit.patch_embed_w->ne[2]; - const int64_t n_img = model.hparams.n_img_embd(); + // Input tensor: pre-computed geometry embeddings after final_proj + norm [D, N_geo, 1] + // The caller pre-computes: final_proj(box_embeds + CLS) → LayerNorm + auto* x = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, N_geo, 1); + ggml_set_name(x, "geom_post_final_proj"); // also used as upload target + ggml_set_input(x); - if (ne[0] != patch_k || ne[1] != n_img || ne[2] != n_img) { - fprintf(stderr, "%s: patch_mulmat input shape mismatch [%lld,%lld,%lld,%lld]\n", - __func__, - (long long) ne[0], (long long) ne[1], (long long) ne[2], (long long) ne[3]); - return false; - } + // Transformer layers (3 layers: self-attn + cross-attn + FFN, pre-norm) + for (int i = 0; i < (int)ge.layers.size(); ++i) { + const auto& ly = ge.layers[i]; - const int64_t n_patch = ne[1] * ne[2] * ne[3]; - const int64_t n_out = model.vit.patch_embed_w->ne[3]; + // 1. Self-attention (pre-norm, pos_enc_at_attn=False) + { + auto* shortcut = x; + auto* xn = sam3_layer_norm(ctx, x, ly.norm1_w, ly.norm1_b); - std::vector input_f16((size_t) (patch_k * n_patch)); - ggml_fp32_to_fp16_row(input_data, input_f16.data(), patch_k * n_patch); + // Q = K = V = norm(x) — no positional encoding at self-attention + auto* Q = ggml_add(ctx, ggml_mul_mat(ctx, ggml_view_2d(ctx, ly.sa_in_proj_w, D, D, ly.sa_in_proj_w->nb[1], 0), xn), + ggml_view_1d(ctx, ly.sa_in_proj_b, D, 0)); + auto* K = ggml_add(ctx, ggml_mul_mat(ctx, ggml_view_2d(ctx, ly.sa_in_proj_w, D, D, ly.sa_in_proj_w->nb[1], D * ly.sa_in_proj_w->nb[1]), xn), + ggml_view_1d(ctx, ly.sa_in_proj_b, D, D * sizeof(float))); + auto* V = ggml_add(ctx, ggml_mul_mat(ctx, ggml_view_2d(ctx, ly.sa_in_proj_w, D, D, ly.sa_in_proj_w->nb[1], 2 * D * ly.sa_in_proj_w->nb[1]), xn), + ggml_view_1d(ctx, ly.sa_in_proj_b, D, 2 * D * sizeof(float))); - std::vector weight_f16((size_t) (patch_k * n_out)); - ggml_backend_tensor_get(model.vit.patch_embed_w, weight_f16.data(), 0, weight_f16.size() * sizeof(ggml_fp16_t)); + const int64_t S = N_geo; + const int64_t HD = D / n_heads; - output_data.resize((size_t) (n_patch * n_out)); - output_ne[0] = n_patch; - output_ne[1] = n_out; - output_ne[2] = 1; - output_ne[3] = 1; + Q = ggml_reshape_4d(ctx, Q, HD, n_heads, S, 1); + Q = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 1, 3)); + K = ggml_reshape_4d(ctx, K, HD, n_heads, S, 1); + K = ggml_cont(ctx, ggml_permute(ctx, K, 0, 2, 1, 3)); + V = ggml_reshape_4d(ctx, V, HD, n_heads, S, 1); + V = ggml_permute(ctx, V, 0, 2, 1, 3); - for (int64_t oc = 0; oc < n_out; ++oc) { - const ggml_fp16_t * w_row = weight_f16.data() + oc * patch_k; - float * dst_row = output_data.data() + oc * n_patch; + float scale = 1.0f / sqrtf((float)HD); + auto* sa_out = ggml_flash_attn_ext(ctx, Q, K, V, nullptr, scale, 0.0f, 0.0f); + sa_out = ggml_reshape_3d(ctx, sa_out, D, S, 1); - for (int64_t p = 0; p < n_patch; ++p) { - const ggml_fp16_t * x_row = input_f16.data() + p * patch_k; + sa_out = ggml_mul_mat(ctx, ly.sa_out_proj_w, sa_out); + sa_out = ggml_add(ctx, sa_out, ly.sa_out_proj_b); - if (use_double_accum) { - double acc = 0.0; - for (int64_t k = 0; k < patch_k; ++k) { - acc += (double) ggml_fp16_to_fp32(x_row[k]) * (double) ggml_fp16_to_fp32(w_row[k]); - } - dst_row[p] = (float) acc; - } else { - float acc = 0.0f; - for (int64_t k = 0; k < patch_k; ++k) { - acc += ggml_fp16_to_fp32(x_row[k]) * ggml_fp16_to_fp32(w_row[k]); - } - dst_row[p] = acc; - } + x = ggml_add(ctx, shortcut, sa_out); } - } - - return true; -} - -bool sam3_test_run_vit_block_linear_host_ref(const sam3_model & model, - int block_idx, - sam3_vit_block_stage stage, - const float * input_data, - const int64_t input_ne[4], - bool use_double_accum, - std::vector & output_data, - int64_t output_ne[4]) { - if (!input_data || block_idx < 0 || block_idx >= (int) model.vit.blocks.size()) { - return false; - } - const sam3_vit_block & blk = model.vit.blocks[(size_t) block_idx]; - - const ggml_tensor * w = nullptr; - const ggml_tensor * b = nullptr; - - switch (stage) { - case SAM3_VIT_BLOCK_STAGE_QKV_PROJ: - w = blk.qkv_w; - b = blk.qkv_b; - break; - case SAM3_VIT_BLOCK_STAGE_ATTN_PROJ: - w = blk.proj_w; - b = blk.proj_b; - break; - case SAM3_VIT_BLOCK_STAGE_MLP_FC1: - w = blk.mlp_fc1_w; - b = blk.mlp_fc1_b; - break; - case SAM3_VIT_BLOCK_STAGE_MLP_FC2: - w = blk.mlp_fc2_w; - b = blk.mlp_fc2_b; - break; - default: - fprintf(stderr, "%s: unsupported stage %d\n", __func__, (int) stage); - return false; - } - - if (!w || !b || w->type != GGML_TYPE_F16 || b->type != GGML_TYPE_F32) { - fprintf(stderr, "%s: unsupported tensor types for stage %d\n", __func__, (int) stage); - return false; - } - - int64_t ne[4]; - sam3_normalize_ne4(input_ne, ne); + // 2. Cross-attention (pre-norm, Q from x, K from img+PE, V from img) + { + auto* shortcut = x; + auto* xn = sam3_layer_norm(ctx, x, ly.norm2_w, ly.norm2_b); - const int64_t k = w->ne[0]; - const int64_t out_dim = w->ne[1]; - const int64_t n_col = ne[1] * ne[2] * ne[3]; + // Q from normalized geometry tokens (no pos) + // K from image features + PE + // V from image features + auto* q_w = ggml_view_2d(ctx, ly.ca_q_w, D, D, ly.ca_q_w->nb[1], 0); + auto* k_w = ggml_view_2d(ctx, ly.ca_q_w, D, D, ly.ca_q_w->nb[1], D * ly.ca_q_w->nb[1]); + auto* v_w = ggml_view_2d(ctx, ly.ca_q_w, D, D, ly.ca_q_w->nb[1], 2 * D * ly.ca_q_w->nb[1]); + auto* q_b = ggml_view_1d(ctx, ly.ca_q_b, D, 0); + auto* k_b = ggml_view_1d(ctx, ly.ca_q_b, D, D * sizeof(float)); + auto* v_b = ggml_view_1d(ctx, ly.ca_q_b, D, 2 * D * sizeof(float)); - if (ne[0] != k) { - fprintf(stderr, "%s: linear input shape mismatch [%lld,%lld,%lld,%lld] for K=%lld\n", - __func__, - (long long) ne[0], (long long) ne[1], (long long) ne[2], (long long) ne[3], - (long long) k); - return false; - } + auto* k_input = ggml_add(ctx, img_feats, img_pe); // pos_enc_at_cross_attn_keys - std::vector input_f16((size_t) (k * n_col)); - ggml_fp32_to_fp16_row(input_data, input_f16.data(), k * n_col); + auto* Q = ggml_add(ctx, ggml_mul_mat(ctx, q_w, xn), q_b); + auto* K = ggml_add(ctx, ggml_mul_mat(ctx, k_w, k_input), k_b); + auto* V = ggml_add(ctx, ggml_mul_mat(ctx, v_w, img_feats), v_b); - std::vector weight_f16((size_t) (k * out_dim)); - ggml_backend_tensor_get(w, weight_f16.data(), 0, weight_f16.size() * sizeof(ggml_fp16_t)); + const int64_t S_q = N_geo; + const int64_t S_kv = img_feats->ne[1]; + const int64_t HD = D / n_heads; - std::vector bias_f32((size_t) out_dim); - ggml_backend_tensor_get(b, bias_f32.data(), 0, bias_f32.size() * sizeof(float)); + Q = ggml_reshape_4d(ctx, Q, HD, n_heads, S_q, 1); + Q = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 1, 3)); + K = ggml_reshape_4d(ctx, K, HD, n_heads, S_kv, 1); + K = ggml_cont(ctx, ggml_permute(ctx, K, 0, 2, 1, 3)); + V = ggml_reshape_4d(ctx, V, HD, n_heads, S_kv, 1); + V = ggml_permute(ctx, V, 0, 2, 1, 3); - output_data.resize((size_t) (out_dim * n_col)); - output_ne[0] = out_dim; - output_ne[1] = ne[1]; - output_ne[2] = ne[2]; - output_ne[3] = ne[3]; + float scale = 1.0f / sqrtf((float)HD); + auto* ca_out = ggml_flash_attn_ext(ctx, Q, K, V, nullptr, scale, 0.0f, 0.0f); + ca_out = ggml_reshape_3d(ctx, ca_out, D, S_q, 1); - for (int64_t col = 0; col < n_col; ++col) { - const ggml_fp16_t * x_col = input_f16.data() + col * k; - float * dst_col = output_data.data() + col * out_dim; + ca_out = ggml_mul_mat(ctx, ly.ca_out_w, ca_out); + ca_out = ggml_add(ctx, ca_out, ly.ca_out_b); - for (int64_t oc = 0; oc < out_dim; ++oc) { - const ggml_fp16_t * w_row = weight_f16.data() + oc * k; + x = ggml_add(ctx, shortcut, ca_out); + } - if (use_double_accum) { - double acc = (double) bias_f32[oc]; - for (int64_t i = 0; i < k; ++i) { - acc += (double) ggml_fp16_to_fp32(w_row[i]) * (double) ggml_fp16_to_fp32(x_col[i]); - } - dst_col[oc] = (float) acc; - } else { - float acc = bias_f32[oc]; - for (int64_t i = 0; i < k; ++i) { - acc += ggml_fp16_to_fp32(w_row[i]) * ggml_fp16_to_fp32(x_col[i]); - } - dst_col[oc] = acc; - } + // 3. FFN (pre-norm, ReLU) + { + auto* shortcut = x; + auto* xn = sam3_layer_norm(ctx, x, ly.norm3_w, ly.norm3_b); + auto* ffn = ggml_mul_mat(ctx, ly.ffn_fc1_w, xn); + ffn = ggml_add(ctx, ffn, ly.ffn_fc1_b); + ffn = ggml_relu(ctx, ffn); + ffn = ggml_mul_mat(ctx, ly.ffn_fc2_w, ffn); + ffn = ggml_add(ctx, ffn, ly.ffn_fc2_b); + x = ggml_add(ctx, shortcut, ffn); } + sam3_name_tensorf(x, "geom_layer%d_out", i); } - return true; -} + // Final encode norm + x = sam3_layer_norm(ctx, x, ge.encode_norm_w, ge.encode_norm_b); + ggml_set_name(x, "geom_output"); + ggml_set_output(x); -static void sam3_normalize_ne4(const int64_t input_ne[4], int64_t ne[4]) { - for (int i = 0; i < 4; ++i) { - ne[i] = (input_ne && input_ne[i] > 0) ? input_ne[i] : 1; - } + sam3_geom_result result; + result.geo_feats = x; + result.n_tokens = N_geo; + return result; } -static struct ggml_tensor * sam3_new_f32_tensor_4d_from_ne(struct ggml_context * ctx, - const int64_t ne[4]) { - return ggml_new_tensor_4d(ctx, GGML_TYPE_F32, ne[0], ne[1], ne[2], ne[3]); -} +// Pre-compute geometry encoder input on CPU: box embeddings + CLS token. +// Reads model weights from GPU, computes embeddings, returns float vector. +// Layout: [D, N_geo] row-major where N_geo = n_boxes + 1 (CLS last). +static std::vector sam3_precompute_geom_input( + const sam3_model& model, + const sam3_pcs_params& params, + const float* img_feats_data, // [D, W, H] ggml-layout backbone features (nullable if no boxes) + int W_feat, int H_feat) { + const auto& ge = model.geom_enc; + const int D = model.hparams.neck_dim; // 256 + const int roi_size = 7; + const int num_pos_feats = D / 2; // 128 -static bool sam3_copy_tensor_to_f32(struct ggml_tensor * t, std::vector & output) { - if (!t) { - return false; + // Collect all exemplar boxes + struct box_info { + float cx, cy, w, h; + int label; + }; + std::vector boxes; + for (const auto& b : params.pos_exemplars) { + // API provides XYXY in original image space — convert to normalized CxCyWH [0,1] + float cx = (b.x0 + b.x1) * 0.5f; + float cy = (b.y0 + b.y1) * 0.5f; + float bw = b.x1 - b.x0; + float bh = b.y1 - b.y0; + boxes.push_back({cx, cy, bw, bh, 0}); // label 0 = positive } - if (!ggml_is_contiguous(t)) { - fprintf(stderr, "%s: tensor '%s' is not contiguous\n", __func__, ggml_get_name(t)); - return false; + for (const auto& b : params.neg_exemplars) { + float cx = (b.x0 + b.x1) * 0.5f; + float cy = (b.y0 + b.y1) * 0.5f; + float bw = b.x1 - b.x0; + float bh = b.y1 - b.y0; + boxes.push_back({cx, cy, bw, bh, 1}); // label 1 = negative } - const int64_t numel = ggml_nelements(t); - output.resize((size_t) numel); + const int n_boxes = (int)boxes.size(); + const int N_geo = n_boxes + 1; - if (t->type == GGML_TYPE_F32) { - ggml_backend_tensor_get(t, output.data(), 0, (size_t) numel * sizeof(float)); - return true; - } + // Read needed weights from GPU + std::vector box_proj_w_data(4 * D), box_proj_b_data(D); + std::vector type_embed_data(D * 2); + std::vector cls_data(D); + std::vector box_pos_proj_w_data(258 * D), box_pos_proj_b_data(D); + std::vector box_pool_proj_w_data(7 * 7 * D * D), box_pool_proj_b_data(D); + std::vector img_pre_norm_w_data(D), img_pre_norm_b_data(D); - if (t->type == GGML_TYPE_F16) { - std::vector tmp((size_t) numel); - ggml_backend_tensor_get(t, tmp.data(), 0, (size_t) numel * sizeof(ggml_fp16_t)); - ggml_fp16_to_fp32_row(tmp.data(), output.data(), numel); - return true; - } + auto read_f32 = [](struct ggml_tensor* t, float* dst, size_t n) { + sam3_read_f32(t, dst, n); + }; - fprintf(stderr, "%s: unsupported tensor type %d for '%s'\n", - __func__, (int) t->type, ggml_get_name(t)); - return false; -} + read_f32(ge.box_proj_w, box_proj_w_data.data(), 4 * D); + read_f32(ge.box_proj_b, box_proj_b_data.data(), D); + read_f32(ge.type_embed, type_embed_data.data(), D * 2); + read_f32(ge.cls_token, cls_data.data(), D); + read_f32(ge.box_pos_proj_w, box_pos_proj_w_data.data(), 258 * D); + read_f32(ge.box_pos_proj_b, box_pos_proj_b_data.data(), D); -static struct ggml_tensor * sam3_build_vit_prefix_stage_from_input(struct ggml_context * ctx, - struct ggml_tensor * input, - const sam3_model & model, - sam3_vit_prefix_stage stage) { - switch (stage) { - case SAM3_VIT_PREFIX_STAGE_PATCH_IM2COL: - return ggml_im2col(ctx, - model.vit.patch_embed_w, - input, - model.vit.patch_embed_w->ne[0], - model.vit.patch_embed_w->ne[1], - 0, 0, 1, 1, true, - model.vit.patch_embed_w->type); - - case SAM3_VIT_PREFIX_STAGE_PATCH_MULMAT_RAW: - case SAM3_VIT_PREFIX_STAGE_PATCH_MULMAT: { - struct ggml_tensor * result = ggml_mul_mat( - ctx, - ggml_reshape_2d(ctx, input, input->ne[0], input->ne[3] * input->ne[2] * input->ne[1]), - ggml_reshape_2d(ctx, - model.vit.patch_embed_w, - model.vit.patch_embed_w->ne[0] * model.vit.patch_embed_w->ne[1] * model.vit.patch_embed_w->ne[2], - model.vit.patch_embed_w->ne[3])); - - if (stage == SAM3_VIT_PREFIX_STAGE_PATCH_MULMAT_RAW) { - return result; - } + if (n_boxes > 0) { + read_f32(ge.box_pool_proj_w, box_pool_proj_w_data.data(), 7 * 7 * D * D); + read_f32(ge.box_pool_proj_b, box_pool_proj_b_data.data(), D); + read_f32(ge.img_pre_norm_w, img_pre_norm_w_data.data(), D); + read_f32(ge.img_pre_norm_b, img_pre_norm_b_data.data(), D); + } - result = ggml_reshape_4d(ctx, result, input->ne[1], input->ne[2], input->ne[3], model.vit.patch_embed_w->ne[3]); - return ggml_cont(ctx, ggml_permute(ctx, result, 0, 1, 3, 2)); - } + // Output: [D * N_geo] in row-major [D, N_geo] ggml order + std::vector out(D * N_geo, 0.0f); - case SAM3_VIT_PREFIX_STAGE_PATCH_EMBED: { - struct ggml_tensor * x = ggml_conv_2d_sk_p0(ctx, model.vit.patch_embed_w, input); - return ggml_cont(ctx, ggml_permute(ctx, x, 1, 2, 0, 3)); - } + // Encode each box + for (int bi = 0; bi < n_boxes; ++bi) { + const auto& box = boxes[bi]; + float embed[256] = {}; - case SAM3_VIT_PREFIX_STAGE_POS_ADD: { - struct ggml_tensor * pos_target = ggml_new_tensor_4d( - ctx, GGML_TYPE_F32, - model.hparams.vit_embed_dim, - model.hparams.n_img_embd(), - model.hparams.n_img_embd(), - 1); - struct ggml_tensor * pos_tiled = ggml_repeat(ctx, model.vit.pos_embed, pos_target); - return ggml_add(ctx, input, pos_tiled); + // 1. Direct projection: Linear(4, D) + // box_proj_w: ggml [4, D] = PyTorch [D, 4] + // out = x @ W^T + b where x=[cx,cy,w,h] + { + float coords[4] = {box.cx, box.cy, box.w, box.h}; + for (int d = 0; d < D; ++d) { + float sum = box_proj_b_data[d]; + for (int j = 0; j < 4; ++j) + sum += coords[j] * box_proj_w_data[j + d * 4]; // ggml: [4, D] stride + embed[d] = sum; + } } - case SAM3_VIT_PREFIX_STAGE_LN_PRE_NORM: - return ggml_norm(ctx, input, 1e-5f); + // 2. ROI Align + Conv2d(D, D, 7) pool projection + if (img_feats_data) { + // Apply img_pre_norm (LayerNorm) to image features before pooling + // For simplicity, we compute LayerNorm per spatial position + // Actually, LayerNorm is applied to the [D] dimension at each spatial location + // But we need the full feature map, so let's normalize in-place copy + const int N_spatial = W_feat * H_feat; + std::vector normed_feats(D * N_spatial); - case SAM3_VIT_PREFIX_STAGE_LN_PRE: { - struct ggml_tensor * x = ggml_norm(ctx, input, 1e-5f); - x = ggml_mul_inplace(ctx, x, model.vit.ln_pre_w); - x = ggml_add_inplace(ctx, x, model.vit.ln_pre_b); - return x; - } - } + for (int s = 0; s < N_spatial; ++s) { + // Compute mean and variance over D dimension + float mean = 0.0f; + for (int d = 0; d < D; ++d) + mean += img_feats_data[d + s * D]; + mean /= D; - return nullptr; -} + float var = 0.0f; + for (int d = 0; d < D; ++d) { + float diff = img_feats_data[d + s * D] - mean; + var += diff * diff; + } + var /= D; -static struct ggml_tensor * sam3_build_vit_attn_core_from_qkv(struct ggml_context * ctx, - struct ggml_tensor * qkv, - const sam3_vit_block & blk, - const sam3_hparams & hp) { - const int E = hp.vit_embed_dim; - const int NH = hp.vit_num_heads; - const int HD = hp.vit_head_dim(); + float inv_std = 1.0f / sqrtf(var + 1e-5f); + for (int d = 0; d < D; ++d) { + normed_feats[d + s * D] = + (img_feats_data[d + s * D] - mean) * inv_std * img_pre_norm_w_data[d] + img_pre_norm_b_data[d]; + } + } - const int64_t W_cur = qkv->ne[1]; - const int64_t H_cur = qkv->ne[2]; - const int64_t B_cur = qkv->ne[3]; + // Convert CxCyWH [0,1] → XYXY in feature grid coordinates + float fx0 = (box.cx - box.w * 0.5f) * (float)W_feat; + float fy0 = (box.cy - box.h * 0.5f) * (float)H_feat; + float fx1 = (box.cx + box.w * 0.5f) * (float)W_feat; + float fy1 = (box.cy + box.h * 0.5f) * (float)H_feat; - struct ggml_tensor * cur = ggml_reshape_4d(ctx, qkv, E, 3, W_cur * H_cur, B_cur); - cur = ggml_cont(ctx, ggml_permute(ctx, cur, 0, 3, 1, 2)); + // ROI Align + std::vector roi_data(D * roi_size * roi_size); + sam3_roi_align_single(normed_feats.data(), D, W_feat, H_feat, + fx0, fy0, fx1, fy1, roi_size, roi_data.data()); - struct ggml_tensor * Q = ggml_view_3d(ctx, cur, E, W_cur * H_cur, B_cur, - cur->nb[1], cur->nb[2], 0); - struct ggml_tensor * K = ggml_view_3d(ctx, cur, E, W_cur * H_cur, B_cur, - cur->nb[1], cur->nb[2], 1 * cur->nb[3]); - struct ggml_tensor * V = ggml_view_3d(ctx, cur, E, W_cur * H_cur, B_cur, - cur->nb[1], cur->nb[2], 2 * cur->nb[3]); + // Conv2d(D, D, 7): kernel [7, 7, D, D] in ggml = [D_out, D_in, kH, kW] in PyTorch + // Since roi_size = 7 = kernel_size, output is [D, 1, 1] + // This is effectively a matrix multiply: out[d_out] = sum over (d_in, kh, kw) + for (int d_out = 0; d_out < D; ++d_out) { + float sum = box_pool_proj_b_data[d_out]; + for (int kh = 0; kh < roi_size; ++kh) { + for (int kw = 0; kw < roi_size; ++kw) { + for (int d_in = 0; d_in < D; ++d_in) { + // ggml weight layout: [kW=7, kH=7, D_in=256, D_out=256] + int w_idx = kw + kh * 7 + d_in * 7 * 7 + d_out * 7 * 7 * D; + // roi_data layout: [C, W, H] = [D_in, kW, kH] + int r_idx = d_in + kw * D + kh * D * roi_size; + sum += box_pool_proj_w_data[w_idx] * roi_data[r_idx]; + } + } + } + embed[d_out] += sum; + } + } - Q = ggml_reshape_4d(ctx, Q, HD, NH, W_cur * H_cur, B_cur); - Q = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 1, 3)); - Q = ggml_reshape_3d(ctx, Q, HD, W_cur * H_cur, NH * B_cur); + // 3. Sinusoidal positional encoding + Linear(258, D) + { + float pos_enc[258]; + sam3_sine_encode_box(pos_enc, box.cx, box.cy, box.w, box.h, + num_pos_feats, 10000); - K = ggml_reshape_4d(ctx, K, HD, NH, W_cur * H_cur, B_cur); - K = ggml_cont(ctx, ggml_permute(ctx, K, 0, 2, 1, 3)); - K = ggml_reshape_3d(ctx, K, HD, W_cur * H_cur, NH * B_cur); + for (int d = 0; d < D; ++d) { + float sum = box_pos_proj_b_data[d]; + for (int j = 0; j < 258; ++j) + sum += pos_enc[j] * box_pos_proj_w_data[j + d * 258]; + embed[d] += sum; + } + } - V = ggml_reshape_4d(ctx, V, HD, NH, W_cur * H_cur, B_cur); - V = ggml_permute(ctx, V, 0, 2, 1, 3); + // 4. Label embedding + { + int label = box.label; + for (int d = 0; d < D; ++d) + embed[d] += type_embed_data[d + label * D]; + } - if (blk.freqs_cis) { - Q = sam3_apply_rope(ctx, Q, blk.freqs_cis); - K = sam3_apply_rope(ctx, K, blk.freqs_cis); + // Store in output: ggml layout [D, N_geo] — embed goes at column bi + for (int d = 0; d < D; ++d) + out[d + bi * D] = embed[d]; } - Q = ggml_reshape_4d(ctx, Q, HD, W_cur * H_cur, NH, B_cur); - K = ggml_reshape_4d(ctx, K, HD, W_cur * H_cur, NH, B_cur); + // CLS token goes at position n_boxes (last) + for (int d = 0; d < D; ++d) + out[d + n_boxes * D] = cls_data[d]; - const float scale = 1.0f / sqrtf((float) HD); - struct ggml_tensor * attn_out = ggml_flash_attn_ext(ctx, Q, K, V, nullptr, scale, 0.0f, 0.0f); + // Apply final_proj (Linear(D,D)) + LayerNorm on CPU + std::vector proj_w(D * D), proj_b(D); + std::vector ln_w(D), ln_b(D); + read_f32(ge.post_proj_w, proj_w.data(), D * D); + read_f32(ge.post_proj_b, proj_b.data(), D); + read_f32(ge.norm_w, ln_w.data(), D); + read_f32(ge.norm_b, ln_b.data(), D); - return ggml_cont(ctx, ggml_reshape_4d(ctx, attn_out, E, W_cur, H_cur, B_cur)); -} + std::vector projected(D * N_geo); + for (int t = 0; t < N_geo; ++t) { + // Linear: y = W @ x + b + // W is stored in PyTorch row-major [D_out, D_in] → ggml [D_in, D_out] + // proj_w[d_in + d_out * D] = W_py[d_out, d_in] + for (int d_out = 0; d_out < D; ++d_out) { + float sum = proj_b[d_out]; + for (int d_in = 0; d_in < D; ++d_in) + sum += out[d_in + t * D] * proj_w[d_in + d_out * D]; + projected[d_out + t * D] = sum; + } + } + if (n_boxes > 0) { + fprintf(stderr, "%s: %d exemplar boxes encoded (%d pos, %d neg)\n", + __func__, n_boxes, + (int)params.pos_exemplars.size(), + (int)params.neg_exemplars.size()); + } + + // LayerNorm over D dimension for each token + for (int t = 0; t < N_geo; ++t) { + float mean = 0.0f; + for (int d = 0; d < D; ++d) + mean += projected[d + t * D]; + mean /= D; -static struct ggml_tensor * sam3_build_vit_block_stage_from_input(struct ggml_context * ctx, - struct ggml_tensor * input, - const sam3_vit_block & blk, - const sam3_hparams & hp, - sam3_vit_block_stage stage) { - switch (stage) { - case SAM3_VIT_BLOCK_STAGE_NORM1: - return sam3_layer_norm(ctx, input, blk.norm1_w, blk.norm1_b); + float var = 0.0f; + for (int d = 0; d < D; ++d) { + float diff = projected[d + t * D] - mean; + var += diff * diff; + } + var /= D; - case SAM3_VIT_BLOCK_STAGE_WINDOW_PART: - return ggml_win_part(ctx, input, hp.vit_window_size); + float inv_std = 1.0f / sqrtf(var + 1e-5f); + for (int d = 0; d < D; ++d) { + float normalized = (projected[d + t * D] - mean) * inv_std; + out[d + t * D] = normalized * ln_w[d] + ln_b[d]; + } + } + return out; +} - case SAM3_VIT_BLOCK_STAGE_QKV_PROJ: - return ggml_add(ctx, ggml_mul_mat(ctx, blk.qkv_w, input), blk.qkv_b); +/***************************************************************************** +** Fusion encoder — graph building (6 layers) +*****************************************************************************/ - case SAM3_VIT_BLOCK_STAGE_ATTN_CORE: - return sam3_build_vit_attn_core_from_qkv(ctx, input, blk, hp); +// Single fusion encoder layer. +// x: [D, N, B] image features (N=5184), prompt: [D, T, B] text/exemplar tokens, pos: [D, N, B] +// Returns: updated x [D, N, B] +static struct ggml_tensor* sam3_fenc_layer_forward( + struct ggml_context* ctx, + const sam3_fenc_layer& ly, + struct ggml_tensor* x, + struct ggml_tensor* prompt, + struct ggml_tensor* pos, + struct ggml_tensor* prompt_attn_bias, + int n_heads) { + // Self-attention: Q/K get positional encoding, V does not + { + auto* shortcut = x; + auto* x_norm = sam3_layer_norm(ctx, x, ly.norm1_w, ly.norm1_b); + auto* q_in = ggml_add(ctx, x_norm, pos); + auto* k_in = ggml_add(ctx, x_norm, pos); - case SAM3_VIT_BLOCK_STAGE_ATTN_PROJ: - return ggml_add(ctx, ggml_mul_mat(ctx, blk.proj_w, input), blk.proj_b); + const int64_t D = x->ne[0]; - case SAM3_VIT_BLOCK_STAGE_WINDOW_UNPART: - return ggml_win_unpart(ctx, input, hp.n_img_embd(), hp.n_img_embd(), hp.vit_window_size); + auto* q_w = ggml_view_2d(ctx, ly.sa_in_proj_w, D, D, ly.sa_in_proj_w->nb[1], 0); + auto* k_w = ggml_view_2d(ctx, ly.sa_in_proj_w, D, D, ly.sa_in_proj_w->nb[1], D * ly.sa_in_proj_w->nb[1]); + auto* v_w = ggml_view_2d(ctx, ly.sa_in_proj_w, D, D, ly.sa_in_proj_w->nb[1], 2 * D * ly.sa_in_proj_w->nb[1]); - case SAM3_VIT_BLOCK_STAGE_NORM2: - return sam3_layer_norm(ctx, input, blk.norm2_w, blk.norm2_b); + auto* q_b = ggml_view_1d(ctx, ly.sa_in_proj_b, D, 0); + auto* k_b = ggml_view_1d(ctx, ly.sa_in_proj_b, D, D * sizeof(float)); + auto* v_b = ggml_view_1d(ctx, ly.sa_in_proj_b, D, 2 * D * sizeof(float)); - case SAM3_VIT_BLOCK_STAGE_MLP_FC1: - return ggml_add(ctx, ggml_mul_mat(ctx, blk.mlp_fc1_w, input), blk.mlp_fc1_b); + auto* Q = ggml_add(ctx, ggml_mul_mat(ctx, q_w, q_in), q_b); + auto* K = ggml_add(ctx, ggml_mul_mat(ctx, k_w, k_in), k_b); + auto* V = ggml_add(ctx, ggml_mul_mat(ctx, v_w, x_norm), v_b); - case SAM3_VIT_BLOCK_STAGE_MLP_GELU: - return ggml_gelu_erf(ctx, input); + const int64_t N = x->ne[1]; + const int64_t B = x->ne[2]; + const int64_t HD = D / n_heads; - case SAM3_VIT_BLOCK_STAGE_MLP_FC2: - return ggml_add(ctx, ggml_mul_mat(ctx, blk.mlp_fc2_w, input), blk.mlp_fc2_b); + Q = ggml_reshape_4d(ctx, Q, HD, n_heads, N, B); + Q = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 1, 3)); + K = ggml_reshape_4d(ctx, K, HD, n_heads, N, B); + K = ggml_cont(ctx, ggml_permute(ctx, K, 0, 2, 1, 3)); + V = ggml_reshape_4d(ctx, V, HD, n_heads, N, B); + V = ggml_permute(ctx, V, 0, 2, 1, 3); - case SAM3_VIT_BLOCK_STAGE_MLP: { - struct ggml_tensor * x = ggml_mul_mat(ctx, blk.mlp_fc1_w, input); - x = ggml_add(ctx, x, blk.mlp_fc1_b); - x = ggml_gelu_erf(ctx, x); - x = ggml_mul_mat(ctx, blk.mlp_fc2_w, x); - x = ggml_add(ctx, x, blk.mlp_fc2_b); - return x; - } - } + float scale = 1.0f / sqrtf((float)HD); + auto* sa_out = ggml_flash_attn_ext(ctx, Q, K, V, nullptr, scale, 0.0f, 0.0f); + sa_out = ggml_reshape_3d(ctx, sa_out, D, N, B); - return nullptr; -} + sa_out = ggml_mul_mat(ctx, ly.sa_out_proj_w, sa_out); + sa_out = ggml_add(ctx, sa_out, ly.sa_out_proj_b); -bool sam3_test_run_vit_block_stage(const sam3_model & model, - int block_idx, - sam3_vit_block_stage stage, - const float * input_data, - const int64_t input_ne[4], - std::vector & output_data, - int64_t output_ne[4], - int n_threads) { - if (!input_data) { - fprintf(stderr, "%s: input_data is null\n", __func__); - return false; - } - if (block_idx < 0 || block_idx >= (int) model.vit.blocks.size()) { - fprintf(stderr, "%s: invalid block_idx=%d\n", __func__, block_idx); - return false; + x = ggml_add(ctx, shortcut, sa_out); } - const auto & hp = model.hparams; - const auto & blk = model.vit.blocks[block_idx]; + // Cross-attention: Q from image features, K/V from prompt tokens. + // ca_q_w stores fused [D, 3*D] weights split as Q-proj, K-proj, V-proj. + { + auto* shortcut = x; + auto* x_norm = sam3_layer_norm(ctx, x, ly.norm2_w, ly.norm2_b); + const int64_t D = x->ne[0]; - int64_t ne[4]; - sam3_normalize_ne4(input_ne, ne); + auto* q_w = ggml_view_2d(ctx, ly.ca_q_w, D, D, ly.ca_q_w->nb[1], 0); + auto* k_w = ggml_view_2d(ctx, ly.ca_q_w, D, D, ly.ca_q_w->nb[1], D * ly.ca_q_w->nb[1]); + auto* v_w = ggml_view_2d(ctx, ly.ca_q_w, D, D, ly.ca_q_w->nb[1], 2 * D * ly.ca_q_w->nb[1]); - switch (stage) { - case SAM3_VIT_BLOCK_STAGE_WINDOW_PART: - case SAM3_VIT_BLOCK_STAGE_WINDOW_UNPART: - if (ne[0] != hp.vit_embed_dim) { - fprintf(stderr, "%s: window stage input ne0=%lld expected %d\n", - __func__, (long long) ne[0], hp.vit_embed_dim); - return false; - } - break; - case SAM3_VIT_BLOCK_STAGE_QKV_PROJ: - if (ne[0] != hp.vit_embed_dim) { - fprintf(stderr, "%s: qkv input ne0=%lld expected %d\n", - __func__, (long long) ne[0], hp.vit_embed_dim); - return false; - } - break; - case SAM3_VIT_BLOCK_STAGE_ATTN_CORE: - if (ne[0] != 3 * hp.vit_embed_dim) { - fprintf(stderr, "%s: attn input ne0=%lld expected %d\n", - __func__, (long long) ne[0], 3 * hp.vit_embed_dim); - return false; - } - break; - case SAM3_VIT_BLOCK_STAGE_MLP_FC1: - if (ne[0] != hp.vit_embed_dim) { - fprintf(stderr, "%s: mlp_fc1 input ne0=%lld expected %d\n", - __func__, (long long) ne[0], hp.vit_embed_dim); - return false; - } - break; - case SAM3_VIT_BLOCK_STAGE_MLP_GELU: - case SAM3_VIT_BLOCK_STAGE_MLP_FC2: - if (ne[0] != hp.vit_mlp_dim) { - fprintf(stderr, "%s: mlp stage input ne0=%lld expected %d\n", - __func__, (long long) ne[0], hp.vit_mlp_dim); - return false; - } - break; - default: - if (ne[0] != hp.vit_embed_dim) { - fprintf(stderr, "%s: stage input ne0=%lld expected %d\n", - __func__, (long long) ne[0], hp.vit_embed_dim); - return false; - } - break; - } + auto* q_b = ggml_view_1d(ctx, ly.ca_q_b, D, 0); + auto* k_b = ggml_view_1d(ctx, ly.ca_q_b, D, D * sizeof(float)); + auto* v_b = ggml_view_1d(ctx, ly.ca_q_b, D, 2 * D * sizeof(float)); - const size_t ctx_size = ggml_tensor_overhead() * 256 + ggml_graph_overhead(); - ggml_init_params params = { - /*.mem_size =*/ ctx_size, - /*.mem_buffer =*/ nullptr, - /*.no_alloc =*/ true, - }; + auto* Q = ggml_add(ctx, ggml_mul_mat(ctx, q_w, x_norm), q_b); + auto* K = ggml_add(ctx, ggml_mul_mat(ctx, k_w, prompt), k_b); + auto* V = ggml_add(ctx, ggml_mul_mat(ctx, v_w, prompt), v_b); - struct ggml_context * ctx = ggml_init(params); - if (!ctx) { - fprintf(stderr, "%s: failed to create ggml context\n", __func__); - return false; - } + const int64_t N_q = x->ne[1]; + const int64_t N_kv = prompt->ne[1]; + const int64_t B = x->ne[2]; + const int64_t HD = D / n_heads; - struct ggml_tensor * input = sam3_new_f32_tensor_4d_from_ne(ctx, ne); - ggml_set_name(input, "vit_block_stage_input"); - ggml_set_input(input); + Q = ggml_reshape_4d(ctx, Q, HD, n_heads, N_q, B); + Q = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 1, 3)); + K = ggml_reshape_4d(ctx, K, HD, n_heads, N_kv, B); + K = ggml_cont(ctx, ggml_permute(ctx, K, 0, 2, 1, 3)); + V = ggml_reshape_4d(ctx, V, HD, n_heads, N_kv, B); + V = ggml_permute(ctx, V, 0, 2, 1, 3); - struct ggml_tensor * output = sam3_build_vit_block_stage_from_input(ctx, input, blk, hp, stage); - if (!output) { - fprintf(stderr, "%s: failed to build stage %d\n", __func__, (int) stage); - ggml_free(ctx); - return false; - } - ggml_set_name(output, "vit_block_stage_output"); - ggml_set_output(output); + auto* ca_mask = sam3_expand_token_attn_bias(ctx, prompt_attn_bias, N_q, n_heads, B); + float scale = 1.0f / sqrtf((float)HD); + auto* ca_out = ggml_flash_attn_ext(ctx, Q, K, V, ca_mask, scale, 0.0f, 0.0f); + ca_out = ggml_reshape_3d(ctx, ca_out, D, N_q, B); - struct ggml_cgraph * graph = ggml_new_graph_custom(ctx, 1024, false); - ggml_build_forward_expand(graph, output); + ca_out = ggml_mul_mat(ctx, ly.ca_out_w, ca_out); + ca_out = ggml_add(ctx, ca_out, ly.ca_out_b); - auto * galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); - if (!ggml_gallocr_reserve(galloc, graph) || !ggml_gallocr_alloc_graph(galloc, graph)) { - fprintf(stderr, "%s: failed to allocate graph for stage %d\n", __func__, (int) stage); - ggml_gallocr_free(galloc); - ggml_free(ctx); - return false; + x = ggml_add(ctx, shortcut, ca_out); } - const size_t input_bytes = (size_t) ne[0] * (size_t) ne[1] * (size_t) ne[2] * (size_t) ne[3] * sizeof(float); - ggml_backend_tensor_set(input, input_data, 0, input_bytes); - sam3_graph_compute(model.backend, graph, n_threads); + { + auto* shortcut = x; + auto* x_norm = sam3_layer_norm(ctx, x, ly.norm3_w, ly.norm3_b); - for (int i = 0; i < 4; ++i) { - output_ne[i] = output->ne[i]; + auto* ffn = ggml_mul_mat(ctx, ly.ffn_fc1_w, x_norm); + ffn = ggml_add(ctx, ffn, ly.ffn_fc1_b); + ffn = ggml_relu(ctx, ffn); + ffn = ggml_mul_mat(ctx, ly.ffn_fc2_w, ffn); + ffn = ggml_add(ctx, ffn, ly.ffn_fc2_b); + + x = ggml_add(ctx, shortcut, ffn); } - const bool ok = sam3_copy_tensor_to_f32(output, output_data); - ggml_gallocr_free(galloc); - ggml_free(ctx); - return ok; + return x; } -// Test-only: encode from pre-preprocessed float data (bypasses C++ resize/normalize). -// chw_data is in [C, H, W] layout, already normalized, with C=3, H=W=img_size. -bool sam3_encode_image_from_preprocessed(sam3_state& state, - const sam3_model& model, - const float* chw_data, - int img_size) { - auto t_start = std::chrono::high_resolution_clock::now(); +// Build full fusion encoder graph (6 layers). +// image_feats: [D, N, B] where N=5184 (72*72), D=256 +// prompt_tokens: [D, T, B] text/exemplar features +// pos_enc: [D, N, B] sinusoidal positional encoding for image features +// Returns: conditioned_features [D, N, B] +static struct ggml_tensor* sam3_build_fenc_graph( + struct ggml_context* ctx, + const sam3_model& model, + struct ggml_tensor* image_feats, + struct ggml_tensor* prompt_tokens, + struct ggml_tensor* pos_enc, + struct ggml_tensor* prompt_attn_bias = nullptr) { const auto& hp = model.hparams; + auto* x = image_feats; - if (img_size != hp.img_size) { - fprintf(stderr, "%s: img_size mismatch: got %d, expected %d\n", - __func__, img_size, hp.img_size); - return false; - } - - fprintf(stderr, "%s: encoding from preprocessed %dx%d\n", __func__, img_size, img_size); - - // EdgeTAM dispatch: build RepViT graph directly from preprocessed data - if (hp.is_edgetam()) { - state.orig_width = img_size; - state.orig_height = img_size; - - const size_t buf_size = ggml_tensor_overhead() * 16384 + ggml_graph_overhead() * 2; - struct ggml_init_params gparams = {buf_size, nullptr, true}; - auto* ctx0 = ggml_init(gparams); - if (!ctx0) return false; + for (int i = 0; i < hp.fenc_layers; ++i) { + x = sam3_fenc_layer_forward(ctx, model.fenc.layers[i], x, prompt_tokens, + pos_enc, prompt_attn_bias, hp.fenc_heads); + sam3_name_tensorf(x, "fenc_layer%d_out", i); + } - auto* inp = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, img_size, img_size, 3, 1); - ggml_set_name(inp, "input_image"); - ggml_set_input(inp); + return x; +} - struct ggml_tensor* stage_outs[4] = {}; - edgetam_build_repvit_graph(ctx0, inp, model, stage_outs); +/***************************************************************************** +** DETR decoder — graph building (6 layers) +*****************************************************************************/ - struct ggml_tensor* fpn_outs[4] = {}; - edgetam_build_fpn_neck_graph(ctx0, stage_outs, model, fpn_outs); +// inverse_sigmoid: log(x / (1 - x)), clamped to avoid inf +// Python reference uses eps=1e-3: x1 = x.clamp(min=eps), x2 = (1-x).clamp(min=eps) +static struct ggml_tensor* sam3_inverse_sigmoid(struct ggml_context* ctx, struct ggml_tensor* x) { + // clamp x to [1e-3, 1-1e-3] to match Python eps=1e-3 + x = ggml_clamp(ctx, x, 1e-3f, 1.0f - 1e-3f); + // log(x / (1 - x)) = log(x) - log(1 - x) + auto* log_x = ggml_log(ctx, x); + // Compute (1 - x) as (-1)*x + 1. We use ggml_scale_bias which takes float + // scalars (no tensor allocation needed, safe in no_alloc contexts). + auto* one_minus = ggml_scale_bias(ctx, x, -1.0f, 1.0f); + auto* log_1mx = ggml_log(ctx, one_minus); + return ggml_sub(ctx, log_x, log_1mx); +} - int n_fpn = 4 - hp.scalp; - for (int i = 0; i < n_fpn; ++i) { - char name[64]; - snprintf(name, sizeof(name), "fpn_out_%d", i); - ggml_set_name(fpn_outs[i], name); - ggml_set_output(fpn_outs[i]); - } +// Box refinement MLP (3 layers: D→D→D→4 with ReLU) +static struct ggml_tensor* sam3_bbox_mlp(struct ggml_context* ctx, + struct ggml_tensor* x, + struct ggml_tensor* w[3], + struct ggml_tensor* b[3]) { + for (int j = 0; j < 3; ++j) { + x = ggml_mul_mat(ctx, w[j], x); + x = ggml_add(ctx, x, b[j]); + if (j < 2) x = ggml_relu(ctx, x); + } + return x; +} - auto* graph = ggml_new_graph_custom(ctx0, 32768, false); - for (int i = 0; i < n_fpn; ++i) - ggml_build_forward_expand(graph, fpn_outs[i]); +// Build sinusoidal positional embedding for 4D reference points in the ggml graph. +// ref_boxes: [4, NQ, B] — (cx, cy, w, h) after sigmoid, B=1 +// sine_dim_t: [1, 64] — pre-computed angle multipliers (2π / 10000^(2i/128)) +// Returns: [512, NQ, B] sinusoidal embedding matching Python gen_sineembed_for_position +static struct ggml_tensor* sam3_build_sine_pos_embed_4d( + struct ggml_context* ctx, + struct ggml_tensor* ref_boxes, // [4, NQ, B] + struct ggml_tensor* sine_dim_t) { // [1, 64] + const int64_t NQ = ref_boxes->ne[1]; - auto* galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); - if (!ggml_gallocr_reserve(galloc, graph) || !ggml_gallocr_alloc_graph(galloc, graph)) { - ggml_gallocr_free(galloc); - ggml_free(ctx0); - return false; - } + // Python output order: [cy, cx, w, h] → coord indices from boxes [cx(0),cy(1),w(2),h(3)] + const int coord_order[4] = {1, 0, 2, 3}; - ggml_backend_tensor_set(inp, chw_data, 0, 3 * img_size * img_size * sizeof(float)); + struct ggml_tensor* coord_embeds[4]; - if (!sam3_graph_compute(model.backend, graph, state.n_threads)) { - ggml_gallocr_free(galloc); - ggml_free(ctx0); - return false; - } + for (int c = 0; c < 4; ++c) { + int ci = coord_order[c]; + // Extract one coordinate: view into ref_boxes [4, NQ, 1] at element ci + auto* coord = ggml_view_2d(ctx, ref_boxes, 1, NQ, + ref_boxes->nb[1], ci * sizeof(float)); // [1, NQ] - // Copy to state (same pattern as edgetam_encode_image) - if (state.buffer) { ggml_backend_buffer_free(state.buffer); state.buffer = nullptr; } - if (state.pe_buf) { ggml_backend_buffer_free(state.pe_buf); state.pe_buf = nullptr; } - if (state.pe_ctx) { ggml_free(state.pe_ctx); state.pe_ctx = nullptr; } - if (state.ctx) { ggml_free(state.ctx); state.ctx = nullptr; } + // Outer product: angles[i, q] = dim_t[i] * coord[q] + // ggml_mul_mat(A=[1,64], B=[1,NQ]) = A^T @ B = [64,1]@[1,NQ] = [64, NQ] + auto* angles = ggml_mul_mat(ctx, sine_dim_t, coord); // [64, NQ] - size_t state_ctx_size = ggml_tensor_overhead() * 32; - struct ggml_init_params sparams = {state_ctx_size, nullptr, true}; - state.ctx = ggml_init(sparams); + auto* sin_vals = ggml_sin(ctx, angles); // [64, NQ] + auto* cos_vals = ggml_cos(ctx, angles); // [64, NQ] - for (int i = 0; i < n_fpn; ++i) { - auto* src = fpn_outs[i]; - state.neck_trk[i] = ggml_new_tensor_4d(state.ctx, GGML_TYPE_F32, - src->ne[0], src->ne[1], src->ne[2], src->ne[3]); - char nm[64]; - snprintf(nm, sizeof(nm), "neck_trk_%d", i); - ggml_set_name(state.neck_trk[i], nm); - } - for (int i = n_fpn; i < 4; ++i) state.neck_trk[i] = nullptr; + // Interleave: [sin_0, cos_0, sin_1, cos_1, ...] + auto* sin_r = ggml_reshape_3d(ctx, sin_vals, 1, 64, NQ); + auto* cos_r = ggml_reshape_3d(ctx, cos_vals, 1, 64, NQ); + auto* interleaved = ggml_concat(ctx, sin_r, cos_r, 0); // [2, 64, NQ] + coord_embeds[c] = ggml_reshape_2d(ctx, interleaved, 128, NQ); + } - state.buffer = ggml_backend_alloc_ctx_tensors(state.ctx, model.backend); - for (int i = 0; i < n_fpn; ++i) { - int64_t nb = ggml_nbytes(state.neck_trk[i]); - std::vector buf(nb); - ggml_backend_tensor_get(fpn_outs[i], buf.data(), 0, nb); - ggml_backend_tensor_set(state.neck_trk[i], buf.data(), 0, nb); - } + // Concatenate all 4 coordinates → [512, NQ] + auto* embed = ggml_concat(ctx, coord_embeds[0], coord_embeds[1], 0); // [256, NQ] + embed = ggml_concat(ctx, embed, coord_embeds[2], 0); // [384, NQ] + embed = ggml_concat(ctx, embed, coord_embeds[3], 0); // [512, NQ] - size_t pe_ctx_size = ggml_tensor_overhead() * 16; - struct ggml_init_params pe_params = {pe_ctx_size, nullptr, true}; - state.pe_ctx = ggml_init(pe_params); - for (int i = 0; i < n_fpn; ++i) { - int H = (int)state.neck_trk[i]->ne[2]; - int W = (int)state.neck_trk[i]->ne[1]; - state.neck_trk_pe[i] = ggml_new_tensor_4d(state.pe_ctx, GGML_TYPE_F32, - hp.neck_dim, W, H, 1); - } - state.pe_buf = ggml_backend_alloc_ctx_tensors(state.pe_ctx, model.backend); - for (int i = 0; i < n_fpn; ++i) { - int H = (int)state.neck_trk[i]->ne[2]; - int W = (int)state.neck_trk[i]->ne[1]; - auto pe = sam3_sinusoidal_pe_2d(H, W, hp.neck_dim); - ggml_backend_tensor_set(state.neck_trk_pe[i], pe.data(), 0, pe.size() * sizeof(float)); - } + return embed; +} - ggml_gallocr_free(galloc); - ggml_free(ctx0); +// Build query positional encoding from reference boxes via sine embed + ref_point_head MLP. +// ref_boxes: [4, NQ, 1] — after sigmoid +// sine_dim_t: [1, 64] +// Returns: [D, NQ+1, 1] (zeros for presence token at index 0, MLP output for object queries) +static struct ggml_tensor* sam3_build_query_pos( + struct ggml_context* ctx, + const sam3_model& model, + struct ggml_tensor* ref_boxes, // [4, NQ, 1] + struct ggml_tensor* sine_dim_t, // [1, 64] + int layer_idx = -1) { + const auto& tensors = model.tensors; + const int64_t NQ = ref_boxes->ne[1]; + const int D = model.hparams.neck_dim; // 256 - auto t_end = std::chrono::high_resolution_clock::now(); - auto ms = std::chrono::duration_cast(t_end - t_start).count(); - fprintf(stderr, "%s: EdgeTAM preprocessed encode done in %lld ms\n", __func__, ms); - return true; + // 1. Sine positional embedding: [512, NQ] + auto* sine_embed = sam3_build_sine_pos_embed_4d(ctx, ref_boxes, sine_dim_t); + if (layer_idx == 0) { + ggml_set_name(sine_embed, "ddec_query_sine_0"); } - // SAM2 dispatch: build a fake sam3_image and use the SAM2 encoder - if (hp.is_sam2()) { - state.orig_width = img_size; - state.orig_height = img_size; + // 2. ref_point_head MLP: 512 → 256 → 256 + // Layer 0: relu(W0 @ sine_embed + b0) + auto* h = ggml_mul_mat(ctx, tensors.at("ddec.ref_point_head.layers.0.weight"), sine_embed); + h = ggml_add(ctx, h, tensors.at("ddec.ref_point_head.layers.0.bias")); + h = ggml_relu(ctx, h); + // Layer 1: W1 @ h + b1 (no activation) + auto* qpos_obj = ggml_mul_mat(ctx, tensors.at("ddec.ref_point_head.layers.1.weight"), h); + qpos_obj = ggml_add(ctx, qpos_obj, tensors.at("ddec.ref_point_head.layers.1.bias")); + // qpos_obj: [D, NQ] + if (layer_idx == 0) { + ggml_set_name(qpos_obj, "ddec_query_pos_0"); + } - // sam2_encode_image_hiera expects an image struct, but we have raw CHW data. - // We need to bypass preprocessing and inject the data directly into the graph. - // Build the Hiera graph, set the input from chw_data (already CHW normalized). + // 3. Reshape to 3D and prepend zeros for presence token + qpos_obj = ggml_reshape_3d(ctx, qpos_obj, D, NQ, 1); + auto* qpos_pres = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, 1, 1); + ggml_set_name(qpos_pres, "ddec_query_pos_pres"); + ggml_set_input(qpos_pres); // zeros — set by caller - const size_t buf_size = ggml_tensor_overhead() * 16384 + ggml_graph_overhead() * 2; - struct ggml_init_params gparams = {buf_size, nullptr, true}; - auto* ctx0 = ggml_init(gparams); - if (!ctx0) return false; + return ggml_concat(ctx, qpos_pres, qpos_obj, 1); // [D, NQ+1, 1] +} - auto* inp = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, img_size, img_size, 3, 1); - ggml_set_name(inp, "input_image"); - ggml_set_input(inp); +// Build box-relative positional bias for DETR cross-attention. +// ref_boxes: [4, N_q, B] — (cx, cy, w, h) in [0,1] +// rpb_coords: [feat_hw] — normalized coords [0/H, 1/H, ..., (H-1)/H] (input tensor) +// Returns: bias tensor [N_kv, N_q+1, n_heads, B] for ggml_flash_attn_ext mask +// +// Python _get_rpb_matrix with boxRPB="log": +// 1. boxes → xyxy +// 2. deltas_x[q,w,:2] = [coord_w - x0, coord_w - x1] +// 3. deltas_y[q,h,:2] = [coord_h - y0, coord_h - y1] +// 4. log transform: sign(d*8) * log2(|d*8|+1) / log2(8) +// 5. MLP: [2] → [256] → [n_heads] +// 6. outer sum: B[h,w] = delta_y[h] + delta_x[w] +static struct ggml_tensor* sam3_compute_box_rpb( + struct ggml_context* ctx, + const sam3_model& model, + struct ggml_tensor* ref_boxes, // [4, N_q, B] + struct ggml_tensor* rpb_coords, // [feat_hw] — pre-filled grid coordinates + int feat_hw, + int layer_idx = -1) { + const int64_t NQ = ref_boxes->ne[1]; + const int NH = model.hparams.ddec_heads; // 8 + const int W = feat_hw; + const int H = feat_hw; + const auto& tensors = model.tensors; - struct ggml_tensor* stage_outs[4] = {}; - sam2_build_hiera_graph(ctx0, inp, model, stage_outs); + // ── 1. Convert cxcywh → xyxy ───────────────────────────────────────── + // ggml_view_2d on strided data is non-contiguous — ggml_scale requires contiguous. + // Use ggml_cont to make each coordinate slice contiguous. + auto* cx = ggml_cont(ctx, ggml_view_2d(ctx, ref_boxes, 1, NQ, ref_boxes->nb[1], 0)); + auto* cy = ggml_cont(ctx, ggml_view_2d(ctx, ref_boxes, 1, NQ, ref_boxes->nb[1], 1 * sizeof(float))); + auto* bw = ggml_cont(ctx, ggml_view_2d(ctx, ref_boxes, 1, NQ, ref_boxes->nb[1], 2 * sizeof(float))); + auto* bh = ggml_cont(ctx, ggml_view_2d(ctx, ref_boxes, 1, NQ, ref_boxes->nb[1], 3 * sizeof(float))); + // x0 = cx - w/2, x1 = cx + w/2 + auto* half_w = ggml_scale(ctx, bw, 0.5f); + auto* half_h = ggml_scale(ctx, bh, 0.5f); + auto* x0 = ggml_sub(ctx, cx, half_w); // [1, NQ] + auto* x1 = ggml_add(ctx, cx, half_w); + auto* y0 = ggml_sub(ctx, cy, half_h); + auto* y1 = ggml_add(ctx, cy, half_h); - struct ggml_tensor* fpn_outs[4] = {}; - sam2_build_fpn_neck_graph(ctx0, stage_outs, model, fpn_outs); + // ── 2. Compute deltas via outer subtract ────────────────────────────── + // coords: [W] → reshape to [W, 1] for outer subtract + auto* cw = ggml_reshape_2d(ctx, rpb_coords, W, 1); // [W, 1] - int n_fpn = 4 - hp.scalp; - for (int i = 0; i < n_fpn; ++i) { - char name[64]; - snprintf(name, sizeof(name), "fpn_out_%d", i); - ggml_set_name(fpn_outs[i], name); - ggml_set_output(fpn_outs[i]); - } + // Outer subtract: delta[w, q] = coord[w] - edge[q] + // Use ggml_mul_mat trick: not applicable for subtraction. + // Instead: repeat coords to [W, NQ], repeat edge to [W, NQ], subtract. + auto* shape_wn = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, W, NQ); - auto* graph = ggml_new_graph_custom(ctx0, 32768, false); - for (int i = 0; i < n_fpn; ++i) - ggml_build_forward_expand(graph, fpn_outs[i]); + auto* cw_rep = ggml_repeat(ctx, cw, shape_wn); // [W, NQ] (each column = coords) + auto* x0_t = ggml_cont(ctx, ggml_transpose(ctx, x0)); // [NQ, 1] + auto* x0_rep = ggml_repeat(ctx, ggml_reshape_2d(ctx, x0_t, 1, NQ), shape_wn); + auto* x1_t = ggml_cont(ctx, ggml_transpose(ctx, x1)); + auto* x1_rep = ggml_repeat(ctx, ggml_reshape_2d(ctx, x1_t, 1, NQ), shape_wn); + auto* y0_t = ggml_cont(ctx, ggml_transpose(ctx, y0)); + auto* y0_rep = ggml_repeat(ctx, ggml_reshape_2d(ctx, y0_t, 1, NQ), shape_wn); + auto* y1_t = ggml_cont(ctx, ggml_transpose(ctx, y1)); + auto* y1_rep = ggml_repeat(ctx, ggml_reshape_2d(ctx, y1_t, 1, NQ), shape_wn); - auto* galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); - if (!ggml_gallocr_reserve(galloc, graph) || !ggml_gallocr_alloc_graph(galloc, graph)) { - ggml_gallocr_free(galloc); - ggml_free(ctx0); - return false; - } + auto* dx0 = ggml_sub(ctx, cw_rep, x0_rep); // [W, NQ] + auto* dx1 = ggml_sub(ctx, cw_rep, x1_rep); + auto* dy0 = ggml_sub(ctx, cw_rep, y0_rep); // reusing coords for H (H==W) + auto* dy1 = ggml_sub(ctx, cw_rep, y1_rep); - // Set input image — CHW data needs to go into ggml's [W, H, C, B] layout - // The input tensor is [img_size, img_size, 3, 1] in ggml = [W, H, C, B] - // CHW data: element (c, h, w) at c*H*W + h*W + w - // ggml data: element at w + h*W_stride... but ggml_conv_2d expects [W, H, C, B] - // which means ne[0]=W, ne[1]=H, ne[2]=C, ne[3]=B - // flat index: w + h*W + c*W*H + b*W*H*C - // CHW: c*H*W + h*W + w → same flat index! So we can just copy directly. - ggml_backend_tensor_set(inp, chw_data, 0, 3 * img_size * img_size * sizeof(float)); + // Stack into [2, W, NQ]: reshape each to [1, W, NQ], concat dim 0 + auto* dx0_r = ggml_reshape_3d(ctx, dx0, 1, W, NQ); + auto* dx1_r = ggml_reshape_3d(ctx, dx1, 1, W, NQ); + auto* deltas_x = ggml_concat(ctx, dx0_r, dx1_r, 0); // [2, W, NQ] + auto* dy0_r = ggml_reshape_3d(ctx, dy0, 1, H, NQ); + auto* dy1_r = ggml_reshape_3d(ctx, dy1, 1, H, NQ); + auto* deltas_y = ggml_concat(ctx, dy0_r, dy1_r, 0); // [2, H, NQ] - // Set positional embedding - { - int pe_H = img_size / 4, pe_W = img_size / 4; - auto pe_data = sam2_compute_pos_embed(model, pe_H, pe_W); - auto* pe_tensor = ggml_graph_get_tensor(graph, "hiera_pos_embed"); - ggml_backend_tensor_set(pe_tensor, pe_data.data(), 0, pe_data.size() * sizeof(float)); - - // Dump PE if requested - const char* dump_dir = getenv("SAM2_DUMP_DIR"); - if (dump_dir) { - char path[512]; - snprintf(path, sizeof(path), "%s/cpp_pos_embed.bin", dump_dir); - FILE* f = fopen(path, "wb"); - if (f) { - fwrite(pe_data.data(), sizeof(float), pe_data.size(), f); - fclose(f); - } - snprintf(path, sizeof(path), "%s/cpp_pos_embed.shape", dump_dir); - f = fopen(path, "w"); - if (f) { - fprintf(f, "%d,%d,%d,%d", hp.hiera_embed_dim, pe_W, pe_H, 1); - fclose(f); - } - fprintf(stderr, " [DUMP] cpp_pos_embed: [%d,%d,%d,1]\n", - hp.hiera_embed_dim, pe_W, pe_H); - } - } + // ── 3. Log transform: sign(d*8) * log2(|d*8|+1) / log2(8) ──────────── + const float scale8 = 8.0f; + const float inv_log2_8 = 1.0f / log2f(8.0f); // = 1/3 - if (ggml_backend_is_cpu(model.backend)) - ggml_backend_cpu_set_n_threads(model.backend, state.n_threads); - if (ggml_backend_graph_compute(model.backend, graph) != GGML_STATUS_SUCCESS) { - ggml_gallocr_free(galloc); - ggml_free(ctx0); - return false; - } + auto rpb_log = [&](struct ggml_tensor* d) -> struct ggml_tensor* { + auto* d8 = ggml_scale(ctx, d, scale8); + auto* sign_d = ggml_sgn(ctx, d8); + auto* abs_d = ggml_abs(ctx, d8); + auto* log_val = ggml_log(ctx, ggml_scale_bias(ctx, abs_d, 1.0f, 1.0f)); + // log2(x) = ln(x) / ln(2) + log_val = ggml_scale(ctx, log_val, 1.0f / logf(2.0f)); + return ggml_mul(ctx, sign_d, ggml_scale(ctx, log_val, inv_log2_8)); + }; - // Dump debug tensors if SAM2_DUMP_DIR is set - { - const char* dump_dir = getenv("SAM2_DUMP_DIR"); - if (dump_dir) { - const char* dbg_names[] = { - "dbg_patch_embed", "dbg_after_pe", - "dbg_blk0_norm1", "dbg_blk0_attn_out", "dbg_blk0_res1", - "dbg_block_0", "dbg_block_1", "dbg_block_2", - "dbg_block_5", "dbg_block_21", - "dbg_fpn_lateral_0", "dbg_fpn_lateral_1", - "dbg_fpn_lateral_2", "dbg_fpn_lateral_3", - }; - for (const char* dn : dbg_names) { - auto* t = ggml_graph_get_tensor(graph, dn); - if (!t) continue; - int64_t nb = ggml_nbytes(t); - std::vector buf(nb); - ggml_backend_tensor_get(t, buf.data(), 0, nb); - char path[512]; - snprintf(path, sizeof(path), "%s/%s.bin", dump_dir, dn); - FILE* f = fopen(path, "wb"); - if (f) { fwrite(buf.data(), 1, nb, f); fclose(f); } - snprintf(path, sizeof(path), "%s/%s.shape", dump_dir, dn); - f = fopen(path, "w"); - if (f) { - fprintf(f, "%lld,%lld,%lld,%lld", - (long long)t->ne[0], (long long)t->ne[1], - (long long)t->ne[2], (long long)t->ne[3]); - fclose(f); - } - fprintf(stderr, " [DUMP] %s: [%lld,%lld,%lld,%lld]\n", dn, - (long long)t->ne[0], (long long)t->ne[1], - (long long)t->ne[2], (long long)t->ne[3]); - } - } - } + deltas_x = rpb_log(deltas_x); // [2, W, NQ] + deltas_y = rpb_log(deltas_y); // [2, H, NQ] - // Copy to state (same as sam2_encode_image_hiera) - if (state.buffer) { ggml_backend_buffer_free(state.buffer); state.buffer = nullptr; } - if (state.pe_buf) { ggml_backend_buffer_free(state.pe_buf); state.pe_buf = nullptr; } - if (state.pe_ctx) { ggml_free(state.pe_ctx); state.pe_ctx = nullptr; } - if (state.ctx) { ggml_free(state.ctx); state.ctx = nullptr; } + // ── 4. MLP: [2, W*NQ] → [NH, W*NQ] ─────────────────────────────────── + // boxRPB_embed_x: MLP(2, 256, 8, 2) = Linear(2→256)+ReLU, Linear(256→8) + // Reshape to [2, W*NQ] so matmul treats each (w, q) pair as a sample + auto rpb_mlp = [&](struct ggml_tensor* d, const char* axis) -> struct ggml_tensor* { + int64_t spatial = d->ne[1]; + int64_t nq = d->ne[2]; + auto* flat = ggml_reshape_2d(ctx, d, 2, spatial * nq); // [2, W*NQ] + auto wn0 = std::string("ddec.boxRPB_embed_") + axis + ".layers.0.weight"; + auto bn0 = std::string("ddec.boxRPB_embed_") + axis + ".layers.0.bias"; + auto wn1 = std::string("ddec.boxRPB_embed_") + axis + ".layers.1.weight"; + auto bn1 = std::string("ddec.boxRPB_embed_") + axis + ".layers.1.bias"; + flat = ggml_mul_mat(ctx, tensors.at(wn0), flat); + flat = ggml_add(ctx, flat, tensors.at(bn0)); + flat = ggml_relu(ctx, flat); + flat = ggml_mul_mat(ctx, tensors.at(wn1), flat); + flat = ggml_add(ctx, flat, tensors.at(bn1)); + // flat: [NH, W*NQ] → reshape to [NH, spatial, NQ] + return ggml_reshape_3d(ctx, flat, NH, spatial, nq); + }; - size_t state_ctx_size = ggml_tensor_overhead() * 32; - struct ggml_init_params sparams = {state_ctx_size, nullptr, true}; - state.ctx = ggml_init(sparams); + auto* rpb_x = rpb_mlp(deltas_x, "x"); // [NH, W, NQ] + auto* rpb_y = rpb_mlp(deltas_y, "y"); // [NH, H, NQ] - for (int i = 0; i < n_fpn; ++i) { - auto* src = fpn_outs[i]; - state.neck_trk[i] = ggml_new_tensor_4d(state.ctx, GGML_TYPE_F32, - src->ne[0], src->ne[1], src->ne[2], src->ne[3]); - } - for (int i = n_fpn; i < 4; ++i) state.neck_trk[i] = nullptr; + // ── 5. Outer sum: B[nh, w, h, q] = rpb_y[nh, h, q] + rpb_x[nh, w, q] ─ + // Reshape for broadcasting: + // rpb_y → [NH, 1, H, NQ] + // rpb_x → [NH, W, 1, NQ] + // + // Keep W in ne[1] so reshaping [NH, W, H, NQ] → [NH, H*W, NQ, 1] + // preserves Python's flatten(H, W) order where W is the fast spatial axis. + auto* rpb_y_4d = ggml_reshape_4d(ctx, rpb_y, NH, 1, H, NQ); + auto* rpb_x_4d = ggml_reshape_4d(ctx, rpb_x, NH, W, 1, NQ); - state.buffer = ggml_backend_alloc_ctx_tensors(state.ctx, model.backend); - for (int i = 0; i < n_fpn; ++i) { - int64_t n_bytes = ggml_nbytes(state.neck_trk[i]); - std::vector buf(n_bytes); - ggml_backend_tensor_get(fpn_outs[i], buf.data(), 0, n_bytes); - ggml_backend_tensor_set(state.neck_trk[i], buf.data(), 0, n_bytes); - } + // ggml_add broadcasts: where one dim is 1, the other is used + auto* rpb_hw = ggml_repeat(ctx, rpb_y_4d, + ggml_new_tensor_4d(ctx, GGML_TYPE_F32, NH, W, H, NQ)); + auto* rpb_hw_x = ggml_repeat(ctx, rpb_x_4d, + ggml_new_tensor_4d(ctx, GGML_TYPE_F32, NH, W, H, NQ)); + auto* rpb = ggml_add(ctx, rpb_hw, rpb_hw_x); // [NH, W, H, NQ] - // Compute sinusoidal PE - size_t pe_ctx_size = ggml_tensor_overhead() * 16; - struct ggml_init_params pe_params = {pe_ctx_size, nullptr, true}; - state.pe_ctx = ggml_init(pe_params); - for (int i = 0; i < n_fpn; ++i) { - int H = (int)state.neck_trk[i]->ne[2]; - int W = (int)state.neck_trk[i]->ne[1]; - state.neck_trk_pe[i] = ggml_new_tensor_4d(state.pe_ctx, GGML_TYPE_F32, hp.neck_dim, W, H, 1); - } - state.pe_buf = ggml_backend_alloc_ctx_tensors(state.pe_ctx, model.backend); - for (int i = 0; i < n_fpn; ++i) { - int H = (int)state.neck_trk[i]->ne[2]; - int W = (int)state.neck_trk[i]->ne[1]; - auto pe = sam3_sinusoidal_pe_2d(H, W, hp.neck_dim); - ggml_backend_tensor_set(state.neck_trk_pe[i], pe.data(), 0, pe.size() * sizeof(float)); - } + // ── 6. Reshape to [H*W, NQ, NH, 1] for flash_attn_ext mask ─────────── + // Current: [NH, W, H, NQ]. Need: [N_kv=H*W, NQ, NH, B=1] + rpb = ggml_reshape_4d(ctx, rpb, NH, H * W, NQ, 1); + rpb = ggml_cont(ctx, ggml_permute(ctx, rpb, 2, 0, 1, 3)); // [H*W, NQ, NH, 1] - ggml_gallocr_free(galloc); - ggml_free(ctx0); + // Prepend zeros for presence token: mask for presence token has no box-relative bias + auto* pres_mask = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, H * W, 1, NH, 1); + ggml_set_name(pres_mask, "rpb_pres_zeros"); + ggml_set_input(pres_mask); // zeros — set by caller - fprintf(stderr, "%s: SAM2 encoding from preprocessed done\n", __func__); - return true; + // [H*W, NQ+1, NH, 1] + auto* full_rpb = ggml_concat(ctx, pres_mask, rpb, 1); + full_rpb = ggml_cont(ctx, full_rpb); + if (layer_idx == 0) { + ggml_set_name(full_rpb, "ddec_rpb_mask_0"); } - state.orig_width = img_size; - state.orig_height = img_size; - - // ── Build computation graph (SAM3 path) ── - const size_t buf_size = ggml_tensor_overhead() * 8192 + ggml_graph_overhead() * 2; - struct ggml_init_params gparams = { - /*.mem_size =*/buf_size, - /*.mem_buffer =*/nullptr, - /*.no_alloc =*/true, - }; - struct ggml_context* ctx0 = ggml_init(gparams); - if (!ctx0) { - fprintf(stderr, "%s: failed to init compute context\n", __func__); - return false; - } + return full_rpb; +} - auto* inp = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, img_size, img_size, 3, 1); - ggml_set_name(inp, "input_image"); - ggml_set_input(inp); +// Single DETR decoder layer. +// queries: [D, N_q, B] where N_q = 201 (200 object queries + 1 presence token) +// query_pos: [D, N_q, B] positional encoding for queries +// enc_feats: [D, N_kv, B] conditioned image features from fusion encoder +// enc_pos: [D, N_kv, B] positional encoding for image features +// text_feats: [D, T, B] text features +// rpb_mask: [N_kv, N_q, n_heads, B] box-relative positional bias (or nullptr) +// Returns: updated queries [D, N_q, B] +static struct ggml_tensor* sam3_ddec_layer_forward( + struct ggml_context* ctx, + const sam3_ddec_layer& ly, + struct ggml_tensor* queries, + struct ggml_tensor* query_pos, + struct ggml_tensor* enc_feats, + struct ggml_tensor* enc_pos, + struct ggml_tensor* text_feats, + int n_heads, + struct ggml_tensor* text_attn_bias = nullptr, + struct ggml_tensor* rpb_mask = nullptr, + int layer_idx = -1) { + const int64_t D = queries->ne[0]; - auto* vit_out = sam3_build_vit_graph(ctx0, inp, model); - ggml_set_name(vit_out, "vit_output"); - ggml_set_output(vit_out); + // Python decoder layer order (all post-norm): + // 1. Self-attention → norm2 (post-norm) + // 2. Text cross-attention (ca_text) → catext_norm (post-norm) + // 3. Image cross-attention (cross_attn) → norm1 (post-norm) + // 4. FFN → norm3 (post-norm) + // + // Norm weight mapping: + // ly.norm2_w = ".norm2.weight" = Python norm2 (post-SA) + // ly.norm3_w = ".norm_ca_text.weight" = Python catext_norm (post-text-CA) + // ly.norm1_w = ".norm1.weight" = Python norm1 (post-image-CA) + // ly.norm4_w = ".norm3.weight" = Python norm3 (post-FFN) - // Mark debug intermediate tensors as outputs so graph allocator preserves them + // 1. Self-attention among queries (post-norm) { - const char* dbg_names[] = { - "dbg_patch_embed", - "dbg_after_pos_embed", - "dbg_ln_pre_norm", - "dbg_ln_pre_scale", - "dbg_after_ln_pre", - "dbg_block_15_norm1", - "dbg_block_15_qkv_proj", - "dbg_block_15_q_split", - "dbg_block_15_k_split", - "dbg_block_15_v_split", - "dbg_block_15_q_heads_base", - "dbg_block_15_k_heads_base", - "dbg_block_15_v_heads_base", - "dbg_block_15_q_heads", - "dbg_block_15_k_heads", - "dbg_block_15_v_flash", - "dbg_block_15_q_rope", - "dbg_block_15_k_rope", - "dbg_block_15_q_flash", - "dbg_block_15_k_flash", - "dbg_block_15_attn_out", - "dbg_block_15_attn_proj", - "dbg_block_15_resid1", - "dbg_block_15_norm2", - "dbg_block_15_mlp", - }; - for (const char* dn : dbg_names) { - auto* dt = ggml_get_tensor(ctx0, dn); - if (dt) ggml_set_output(dt); - } - for (int i = 0; i < (int)model.hparams.vit_depth; ++i) { - char bn[64]; - snprintf(bn, sizeof(bn), "dbg_block_%d_out", i); - auto* dt = ggml_get_tensor(ctx0, bn); - if (dt) ggml_set_output(dt); - } - } + // Q = K = queries + query_pos, V = queries (no pos) + auto* q_in = ggml_add(ctx, queries, query_pos); - struct ggml_tensor* neck_det_out[4] = {}; - struct ggml_tensor* neck_trk_out[4]; - if (!model.hparams.visual_only) { - sam3_build_neck_graph(ctx0, vit_out, model.neck_det, neck_det_out); - } - sam3_build_neck_graph(ctx0, vit_out, model.neck_trk, neck_trk_out); + auto* q_w = ggml_view_2d(ctx, ly.sa_in_proj_w, D, D, ly.sa_in_proj_w->nb[1], 0); + auto* k_w = ggml_view_2d(ctx, ly.sa_in_proj_w, D, D, ly.sa_in_proj_w->nb[1], D * ly.sa_in_proj_w->nb[1]); + auto* v_w = ggml_view_2d(ctx, ly.sa_in_proj_w, D, D, ly.sa_in_proj_w->nb[1], 2 * D * ly.sa_in_proj_w->nb[1]); + auto* q_b = ggml_view_1d(ctx, ly.sa_in_proj_b, D, 0); + auto* k_b = ggml_view_1d(ctx, ly.sa_in_proj_b, D, D * sizeof(float)); + auto* v_b = ggml_view_1d(ctx, ly.sa_in_proj_b, D, 2 * D * sizeof(float)); - for (int i = 0; i < 4; ++i) { - char name[64]; - if (!model.hparams.visual_only) { - snprintf(name, sizeof(name), "neck_det_%d", i); - ggml_set_name(neck_det_out[i], name); - ggml_set_output(neck_det_out[i]); - } - snprintf(name, sizeof(name), "neck_trk_%d", i); - ggml_set_name(neck_trk_out[i], name); - ggml_set_output(neck_trk_out[i]); - } + const int64_t N = queries->ne[1]; + const int64_t B = queries->ne[2]; + const int64_t HD = D / n_heads; - struct ggml_cgraph* graph = ggml_new_graph_custom(ctx0, 16384, false); - for (int i = 0; i < 4; ++i) { - if (!model.hparams.visual_only) { - ggml_build_forward_expand(graph, neck_det_out[i]); - } - ggml_build_forward_expand(graph, neck_trk_out[i]); - } + auto* Q = ggml_add(ctx, ggml_mul_mat(ctx, q_w, q_in), q_b); + auto* K = ggml_add(ctx, ggml_mul_mat(ctx, k_w, q_in), k_b); + auto* V = ggml_add(ctx, ggml_mul_mat(ctx, v_w, queries), v_b); - auto* galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); + Q = ggml_reshape_4d(ctx, Q, HD, n_heads, N, B); + Q = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 1, 3)); + K = ggml_reshape_4d(ctx, K, HD, n_heads, N, B); + K = ggml_cont(ctx, ggml_permute(ctx, K, 0, 2, 1, 3)); + V = ggml_reshape_4d(ctx, V, HD, n_heads, N, B); + V = ggml_permute(ctx, V, 0, 2, 1, 3); - if (!ggml_gallocr_reserve(galloc, graph)) { - fprintf(stderr, "%s: failed to reserve graph memory\n", __func__); - ggml_gallocr_free(galloc); - ggml_free(ctx0); - return false; - } + float scale = 1.0f / sqrtf((float)HD); + auto* sa_out = ggml_flash_attn_ext(ctx, Q, K, V, nullptr, scale, 0.0f, 0.0f); + sa_out = ggml_reshape_3d(ctx, sa_out, D, N, B); + sa_out = ggml_mul_mat(ctx, ly.sa_out_proj_w, sa_out); + sa_out = ggml_add(ctx, sa_out, ly.sa_out_proj_b); - if (!ggml_gallocr_alloc_graph(galloc, graph)) { - fprintf(stderr, "%s: failed to allocate graph\n", __func__); - ggml_gallocr_free(galloc); - ggml_free(ctx0); - return false; + queries = ggml_add(ctx, queries, sa_out); + // Post-norm: norm2 (Python's post-SA norm) + queries = sam3_layer_norm(ctx, queries, ly.norm2_w, ly.norm2_b); + if (layer_idx == 0) { + ggml_set_name(queries, "ddec_layer0_after_sa"); + } } - // Upload preprocessed data directly (already in CHW = ggml [W, H, C, B] layout) - const size_t data_bytes = (size_t)3 * img_size * img_size * sizeof(float); - ggml_backend_tensor_set(inp, chw_data, 0, data_bytes); - - // Compute + // 2. Cross-attention to text tokens (post-norm) { - auto t0 = std::chrono::high_resolution_clock::now(); - if (!sam3_graph_compute(model.backend, graph, state.n_threads)) { - ggml_gallocr_free(galloc); - ggml_free(ctx0); - return false; - } - auto t1 = std::chrono::high_resolution_clock::now(); - double ms = std::chrono::duration(t1 - t0).count(); - fprintf(stderr, "%s: graph computed in %.1f ms (%d threads)\n", - __func__, ms, state.n_threads); - } + // Q = queries + query_pos, K = V = text_feats + auto* q_in = ggml_add(ctx, queries, query_pos); - // Cache results in state - if (state.galloc) ggml_gallocr_free(state.galloc); - if (state.ctx) ggml_free(state.ctx); + auto* q_w = ggml_view_2d(ctx, ly.ca_text_q_w, D, D, ly.ca_text_q_w->nb[1], 0); + auto* k_w = ggml_view_2d(ctx, ly.ca_text_q_w, D, D, ly.ca_text_q_w->nb[1], D * ly.ca_text_q_w->nb[1]); + auto* v_w = ggml_view_2d(ctx, ly.ca_text_q_w, D, D, ly.ca_text_q_w->nb[1], 2 * D * ly.ca_text_q_w->nb[1]); + auto* q_b = ggml_view_1d(ctx, ly.ca_text_q_b, D, 0); + auto* k_b = ggml_view_1d(ctx, ly.ca_text_q_b, D, D * sizeof(float)); + auto* v_b = ggml_view_1d(ctx, ly.ca_text_q_b, D, 2 * D * sizeof(float)); - state.ctx = ctx0; - state.galloc = galloc; - state.backend = model.backend; - state.vit_output = vit_out; + const int64_t N_q = queries->ne[1]; + const int64_t N_kv = text_feats->ne[1]; + const int64_t B = queries->ne[2]; + const int64_t HD = D / n_heads; - for (int i = 0; i < 4; ++i) { - state.neck_det[i] = model.hparams.visual_only ? nullptr : neck_det_out[i]; - state.neck_trk[i] = neck_trk_out[i]; - } + auto* Q = ggml_add(ctx, ggml_mul_mat(ctx, q_w, q_in), q_b); + auto* K = ggml_add(ctx, ggml_mul_mat(ctx, k_w, text_feats), k_b); + auto* V = ggml_add(ctx, ggml_mul_mat(ctx, v_w, text_feats), v_b); - // Compute sinusoidal PEs - { - const int neck_dim = hp.neck_dim; - const int scale_sizes[4] = { - hp.n_img_embd() * 4, - hp.n_img_embd() * 2, - hp.n_img_embd(), - hp.n_img_embd() / 2, - }; + Q = ggml_reshape_4d(ctx, Q, HD, n_heads, N_q, B); + Q = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 1, 3)); + K = ggml_reshape_4d(ctx, K, HD, n_heads, N_kv, B); + K = ggml_cont(ctx, ggml_permute(ctx, K, 0, 2, 1, 3)); + V = ggml_reshape_4d(ctx, V, HD, n_heads, N_kv, B); + V = ggml_permute(ctx, V, 0, 2, 1, 3); - if (state.pe_buf) { - ggml_backend_buffer_free(state.pe_buf); - state.pe_buf = nullptr; - } - if (state.pe_ctx) { - ggml_free(state.pe_ctx); - state.pe_ctx = nullptr; + auto* text_mask = sam3_expand_token_attn_bias(ctx, text_attn_bias, N_q, n_heads, B); + float scale = 1.0f / sqrtf((float)HD); + auto* ca_out = ggml_flash_attn_ext(ctx, Q, K, V, text_mask, scale, 0.0f, 0.0f); + ca_out = ggml_reshape_3d(ctx, ca_out, D, N_q, B); + ca_out = ggml_mul_mat(ctx, ly.ca_text_out_w, ca_out); + ca_out = ggml_add(ctx, ca_out, ly.ca_text_out_b); + + queries = ggml_add(ctx, queries, ca_out); + // Post-norm: catext_norm (Python's post-text-CA norm) + queries = sam3_layer_norm(ctx, queries, ly.norm3_w, ly.norm3_b); + if (layer_idx == 0) { + ggml_set_name(queries, "ddec_layer0_after_text_ca"); } + } - struct ggml_init_params pe_params = { - /*.mem_size =*/ggml_tensor_overhead() * 4 + 256, - /*.mem_buffer =*/nullptr, - /*.no_alloc =*/true, - }; - state.pe_ctx = ggml_init(pe_params); + // 3. Cross-attention to conditioned image features (post-norm) + { + // Q = queries + query_pos, K = enc_feats + enc_pos, V = enc_feats + auto* q_in = ggml_add(ctx, queries, query_pos); + auto* k_in = ggml_add(ctx, enc_feats, enc_pos); - struct ggml_tensor* pe_tensors[4]; - for (int i = 0; i < 4; ++i) { - const int S = scale_sizes[i]; - pe_tensors[i] = ggml_new_tensor_4d(state.pe_ctx, GGML_TYPE_F32, neck_dim, S, S, 1); - char name[64]; - snprintf(name, sizeof(name), "pe_%d", i); - ggml_set_name(pe_tensors[i], name); - } + auto* q_w = ggml_view_2d(ctx, ly.ca_q_w, D, D, ly.ca_q_w->nb[1], 0); + auto* k_w = ggml_view_2d(ctx, ly.ca_q_w, D, D, ly.ca_q_w->nb[1], D * ly.ca_q_w->nb[1]); + auto* v_w = ggml_view_2d(ctx, ly.ca_q_w, D, D, ly.ca_q_w->nb[1], 2 * D * ly.ca_q_w->nb[1]); + auto* q_b = ggml_view_1d(ctx, ly.ca_q_b, D, 0); + auto* k_b = ggml_view_1d(ctx, ly.ca_q_b, D, D * sizeof(float)); + auto* v_b = ggml_view_1d(ctx, ly.ca_q_b, D, 2 * D * sizeof(float)); - state.pe_buf = ggml_backend_alloc_ctx_tensors(state.pe_ctx, model.backend); - if (!state.pe_buf) { - fprintf(stderr, "%s: failed to allocate PE buffer\n", __func__); + const int64_t N_q = queries->ne[1]; + const int64_t N_kv = enc_feats->ne[1]; + const int64_t B = queries->ne[2]; + const int64_t HD = D / n_heads; + + auto* Q = ggml_add(ctx, ggml_mul_mat(ctx, q_w, q_in), q_b); + auto* K = ggml_add(ctx, ggml_mul_mat(ctx, k_w, k_in), k_b); + auto* V = ggml_add(ctx, ggml_mul_mat(ctx, v_w, enc_feats), v_b); + + Q = ggml_reshape_4d(ctx, Q, HD, n_heads, N_q, B); + Q = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 1, 3)); + K = ggml_reshape_4d(ctx, K, HD, n_heads, N_kv, B); + K = ggml_cont(ctx, ggml_permute(ctx, K, 0, 2, 1, 3)); + V = ggml_reshape_4d(ctx, V, HD, n_heads, N_kv, B); + V = ggml_permute(ctx, V, 0, 2, 1, 3); + + float scale = 1.0f / sqrtf((float)HD); + struct ggml_tensor* ca_out = nullptr; + + if (rpb_mask) { + // Keep the box-relative positional bias in fp32. The CPU flash-attn + // kernel reads mask storage as fp16, which is good enough for token + // padding masks but introduces avoidable drift here. + auto* kq = ggml_mul_mat(ctx, K, Q); // [N_kv, N_q, NH, B] + kq = ggml_soft_max_ext(ctx, kq, rpb_mask, scale, 0.0f); // [N_kv, N_q, NH, B] + + auto* v_t = ggml_cont(ctx, ggml_transpose(ctx, V)); // [N_kv, HD, NH, B] + ca_out = ggml_mul_mat(ctx, v_t, kq); // [HD, N_q, NH, B] + ca_out = ggml_cont(ctx, ggml_permute(ctx, ca_out, 0, 2, 1, 3)); } else { - for (int i = 0; i < 4; ++i) { - const int S = scale_sizes[i]; - auto pe_data = sam3_sinusoidal_pe_2d(S, S, neck_dim); - ggml_backend_tensor_set(pe_tensors[i], pe_data.data(), 0, pe_data.size() * sizeof(float)); - state.neck_det_pe[i] = pe_tensors[i]; - state.neck_trk_pe[i] = pe_tensors[i]; - } + ca_out = ggml_flash_attn_ext(ctx, Q, K, V, nullptr, scale, 0.0f, 0.0f); + } + + ca_out = ggml_reshape_3d(ctx, ca_out, D, N_q, B); + ca_out = ggml_mul_mat(ctx, ly.ca_out_w, ca_out); + ca_out = ggml_add(ctx, ca_out, ly.ca_out_b); + + queries = ggml_add(ctx, queries, ca_out); + // Post-norm: norm1 (Python's post-image-CA norm) + queries = sam3_layer_norm(ctx, queries, ly.norm1_w, ly.norm1_b); + if (layer_idx == 0) { + ggml_set_name(queries, "ddec_layer0_after_img_ca"); } } - state.pe_cache_valid = false; + // 4. FFN (post-norm, ReLU) + { + auto* ffn = ggml_mul_mat(ctx, ly.ffn_fc1_w, queries); + ffn = ggml_add(ctx, ffn, ly.ffn_fc1_b); + ffn = ggml_relu(ctx, ffn); + ffn = ggml_mul_mat(ctx, ly.ffn_fc2_w, ffn); + ffn = ggml_add(ctx, ffn, ly.ffn_fc2_b); - auto t_end = std::chrono::high_resolution_clock::now(); - double total_ms = std::chrono::duration(t_end - t_start).count(); - fprintf(stderr, "%s: encoded successfully in %.1f ms\n", __func__, total_ms); - return true; -} + queries = ggml_add(ctx, queries, ffn); + // Post-norm: norm3 (Python's post-FFN norm) + queries = sam3_layer_norm(ctx, queries, ly.norm4_w, ly.norm4_b); + if (layer_idx == 0) { + ggml_set_name(queries, "ddec_layer0_full_out"); + } + } -/***************************************************************************** -** Multi-head attention helper (used by fusion encoder, DETR decoder, seg head) -*****************************************************************************/ + return queries; +} -// Standard multi-head attention with fused in_proj. -// q_in, k_in, v_in: [D, N, B] (if fused_qkv, only q_in is used and contains QKV stacked) -// in_proj_w: [D, 3*D] (fused Q/K/V projection) -// in_proj_b: [3*D] -// out_proj_w: [D, D], out_proj_b: [D] -// n_heads: number of attention heads -// Returns: [D, N_q, B] +// DotProductScoring: classify queries against text features via dot product. // -// If separate_kv is true, q_in/k_in/v_in are already separate (no fused proj needed). -// The in_proj is applied to form Q from q_in, and K/V from the concatenated k/v source. -static struct ggml_tensor* sam3_multihead_attn_fused( +// Python reference (DotProductScoring.forward): +// 1. prompt_mlp(prompt) → residual MLP + LN on text features +// 2. mean_pool_text(result, prompt_mask) → pooled [BS, D] (only valid tokens) +// 3. prompt_proj(pooled) → [BS, D] +// 4. hs_proj(hs) → [num_layer, BS, N_q, D] +// 5. matmul(proj_hs, proj_pooled.unsqueeze(-1)) → dot product → [num_layer, BS, N_q, 1] +// 6. scale by 1/sqrt(D) +// 7. clamp to [-12, 12] +// +// query_outputs: [D, N_q, B] — the 200 object query outputs +// text_features: [D, T, B] — text encoder output (already through resizer) +// text_valid_mask: [T, 1, B] — 1.0 for valid tokens, 0.0 for padding (or nullptr for all-valid) +// Returns: class_scores [N_q, B] (one score per query per batch) +static struct ggml_tensor* sam3_dot_product_scoring( struct ggml_context* ctx, - struct ggml_tensor* q_in, // [D, N_q, B] - struct ggml_tensor* kv_in, // [D, N_kv, B] (can be same as q_in for self-attn) - struct ggml_tensor* in_proj_w, // [D, 3*D] — fused QKV weights - struct ggml_tensor* in_proj_b, // [3*D] - struct ggml_tensor* out_proj_w, // [D, D] - struct ggml_tensor* out_proj_b, // [D] - int n_heads, - struct ggml_tensor* attn_mask = nullptr) // [N_kv, N_q] or nullptr + const sam3_model& model, + struct ggml_tensor* query_outputs, // [D, N_q, B] + struct ggml_tensor* text_features, // [D, T, B] + struct ggml_tensor* text_valid_mask) // [T, 1, B] or nullptr { - const int64_t D = q_in->ne[0]; // 256 - const int64_t N_q = q_in->ne[1]; - const int64_t B = q_in->ne[2]; - const int64_t N_kv = kv_in->ne[1]; - const int64_t HD = D / n_heads; + const auto& tensors = model.tensors; + const int64_t D = query_outputs->ne[0]; // 256 + const int64_t T = text_features->ne[1]; + const int64_t B = text_features->ne[2]; - auto* q_w = ggml_view_2d(ctx, in_proj_w, D, D, in_proj_w->nb[1], 0); - auto* k_w = ggml_view_2d(ctx, in_proj_w, D, D, in_proj_w->nb[1], D * in_proj_w->nb[1]); - auto* v_w = ggml_view_2d(ctx, in_proj_w, D, D, in_proj_w->nb[1], 2 * D * in_proj_w->nb[1]); + // Step 1: Apply prompt_mlp on text features (residual MLP + LayerNorm) + auto* text_mlp = text_features; // [D, T, B] + auto* orig_text = text_features; + text_mlp = ggml_mul_mat(ctx, tensors.at("scoring.prompt_mlp.layers.0.weight"), text_mlp); + text_mlp = ggml_add(ctx, text_mlp, tensors.at("scoring.prompt_mlp.layers.0.bias")); + text_mlp = ggml_relu(ctx, text_mlp); + text_mlp = ggml_mul_mat(ctx, tensors.at("scoring.prompt_mlp.layers.1.weight"), text_mlp); + text_mlp = ggml_add(ctx, text_mlp, tensors.at("scoring.prompt_mlp.layers.1.bias")); + text_mlp = ggml_add(ctx, text_mlp, orig_text); + text_mlp = sam3_layer_norm(ctx, text_mlp, + tensors.at("scoring.prompt_mlp.out_norm.weight"), + tensors.at("scoring.prompt_mlp.out_norm.bias")); + // text_mlp: [D, T, B] + ggml_set_name(text_mlp, "scoring_prompt_mlp_out"); - auto* q_b = ggml_view_1d(ctx, in_proj_b, D, 0); - auto* k_b = ggml_view_1d(ctx, in_proj_b, D, D * sizeof(float)); - auto* v_b = ggml_view_1d(ctx, in_proj_b, D, 2 * D * sizeof(float)); + // Step 2: Mean-pool over valid text tokens → [D, 1, B] + // Python: pooled = (prompt * is_valid).sum(0) / num_valid + struct ggml_tensor* text_pooled; + if (text_valid_mask) { + // Permute to [T, D, B] for masked pooling + auto* tp = ggml_cont(ctx, ggml_permute(ctx, text_mlp, 1, 0, 2, 3)); // [T, D, B] + // Mask: text_valid_mask is [T, 1, B] — broadcast multiply zeros out padding + tp = ggml_mul(ctx, tp, text_valid_mask); // [T, D, B] with padding zeroed + // Sum over T dimension: pool_1d with SUM kernel=T + // ggml_pool_1d AVG divides by T; we want SUM then divide by n_valid. + // Use AVG and then scale by T/n_valid? Or use a manual approach. + // Simpler: sum via pool_1d with AVG, then scale by T/n_valid. + // But n_valid is dynamic. Instead: sum = mean * T, then divide by n_valid. + // We pass n_valid as part of the mask: text_valid_mask sums to n_valid. + // pool_1d(masked, AVG, T, T, 0) = sum(masked) / T. Multiply by T → sum(masked). + // Then divide by n_valid. But n_valid is a scalar we know CPU-side. + // For simplicity: compute AVG over ALL T positions (with padding zeroed out). + // This gives sum(valid) / T. To get sum(valid) / n_valid, scale by T / n_valid. + // We embed the scale factor into the mask: mask = (T / n_valid) for valid, 0 for pad. + // Then AVG(mask * features) = sum(valid * T/n_valid) / T = sum(valid) / n_valid. ✓ + // Caller should set mask values to T/n_valid for valid tokens, 0 for padding. + auto* pooled_t = ggml_pool_1d(ctx, tp, GGML_OP_POOL_AVG, (int)T, (int)T, 0); + text_pooled = ggml_cont(ctx, ggml_permute(ctx, pooled_t, 1, 0, 2, 3)); // [D, 1, B] + } else { + // All tokens valid — simple mean + auto* tp = ggml_cont(ctx, ggml_permute(ctx, text_mlp, 1, 0, 2, 3)); + auto* pooled_t = ggml_pool_1d(ctx, tp, GGML_OP_POOL_AVG, (int)T, (int)T, 0); + text_pooled = ggml_cont(ctx, ggml_permute(ctx, pooled_t, 1, 0, 2, 3)); + } + ggml_set_name(text_pooled, "scoring_pooled"); - auto* Q = ggml_add(ctx, ggml_mul_mat(ctx, q_w, q_in), q_b); - auto* K = ggml_add(ctx, ggml_mul_mat(ctx, k_w, kv_in), k_b); - auto* V = ggml_add(ctx, ggml_mul_mat(ctx, v_w, kv_in), v_b); + // Step 3: Project pooled prompt through prompt_proj: D→D + auto* proj_pooled = ggml_mul_mat(ctx, tensors.at("scoring.prompt_proj.weight"), text_pooled); + proj_pooled = ggml_add(ctx, proj_pooled, tensors.at("scoring.prompt_proj.bias")); + // proj_pooled: [D, 1, B] + ggml_set_name(proj_pooled, "scoring_proj_pooled"); - Q = ggml_reshape_4d(ctx, Q, HD, n_heads, N_q, B); - Q = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 1, 3)); // [HD, N_q, NH, B] + // Step 4: Project queries through hs_proj: D→D + auto* proj_hs = ggml_mul_mat(ctx, tensors.at("scoring.hs_proj.weight"), query_outputs); + proj_hs = ggml_add(ctx, proj_hs, tensors.at("scoring.hs_proj.bias")); + // proj_hs: [D, N_q, B] + ggml_set_name(proj_hs, "scoring_proj_hs"); - K = ggml_reshape_4d(ctx, K, HD, n_heads, N_kv, B); - K = ggml_cont(ctx, ggml_permute(ctx, K, 0, 2, 1, 3)); // [HD, N_kv, NH, B] + // Step 5: Dot product — for each query, dot with pooled prompt + // matmul(proj_hs, proj_pooled.unsqueeze(-1)) in Python = batched vector-matrix multiply + // ggml_mul_mat(A, B) = A^T @ B + // With A = proj_pooled [D, 1, B], B = proj_hs [D, N_q, B]: + // result = [1, N_q, B] — each element is dot product of query with pooled prompt + auto* scores = ggml_mul_mat(ctx, proj_pooled, proj_hs); // [1, N_q, B] - V = ggml_reshape_4d(ctx, V, HD, n_heads, N_kv, B); - V = ggml_permute(ctx, V, 0, 2, 1, 3); // [HD, N_kv, NH, B] non-contiguous; flash_attn uses strides + // Step 6: Scale by 1/sqrt(D) + float scale = 1.0f / sqrtf((float)D); + scores = ggml_scale(ctx, scores, scale); - float scale = 1.0f / sqrtf((float)HD); - auto* attn_out = ggml_flash_attn_ext(ctx, Q, K, V, attn_mask, scale, 0.0f, 0.0f); + // Step 7: Clamp to [-12, 12] + scores = ggml_clamp(ctx, scores, -12.0f, 12.0f); - auto* merged = ggml_reshape_3d(ctx, attn_out, D, N_q, B); - merged = ggml_mul_mat(ctx, out_proj_w, merged); - merged = ggml_add(ctx, merged, out_proj_b); + // Reshape to [N_q, B] + const int64_t N_q = query_outputs->ne[1]; + scores = ggml_reshape_2d(ctx, scores, N_q, B); + ggml_set_name(scores, "scoring_class_scores"); - return merged; + return scores; } -static struct ggml_tensor* sam3_expand_token_attn_bias( - struct ggml_context* ctx, - struct ggml_tensor* token_bias, // [T, 1, B] or nullptr - int64_t n_q, - int n_heads, - int64_t batch) { - if (!token_bias) { - return nullptr; - } +// Build full DETR decoder graph. +// enc_feats: [D, N_kv, B] conditioned features from fusion encoder (N_kv=5184) +// enc_pos: [D, N_kv, B] positional encoding +// text_feats: [D, T, B] text features +// Returns struct with: +// queries: [D, 201, B] (all query outputs including presence token) +// pred_boxes: [4, 200, B] (cx, cy, w, h in [0,1]) +// class_scores: [200, B] +// presence_score: [1, B] +struct sam3_ddec_output { + struct ggml_tensor* queries; // [D, 201, B] + struct ggml_tensor* presence_feats; // [D, 1, B] pre-decoder-norm presence token + struct ggml_tensor* pred_boxes; // [4, 200, B] + struct ggml_tensor* class_scores; // [200, B] + struct ggml_tensor* presence_score; // [1, B] +}; - const int64_t n_kv = token_bias->ne[0]; - auto* bias_4d = ggml_reshape_4d(ctx, token_bias, n_kv, 1, 1, batch); - auto* full_bias = ggml_repeat( - ctx, - bias_4d, - ggml_new_tensor_4d(ctx, GGML_TYPE_F32, n_kv, n_q, n_heads, batch)); +static sam3_ddec_output sam3_build_ddec_graph( + struct ggml_context* ctx, + const sam3_model& model, + struct ggml_tensor* enc_feats, // [D, N_kv, B] + struct ggml_tensor* enc_pos, // [D, N_kv, B] + struct ggml_tensor* text_feats, // [D, T, B] + struct ggml_tensor* sine_dim_t, // [1, 64] — pre-computed angle multipliers + struct ggml_tensor* rpb_coords, // [feat_hw] — normalized grid coords (or nullptr) + struct ggml_tensor* text_attn_bias = nullptr, // [T, 1, B] additive text padding bias + struct ggml_tensor* text_valid_mask = nullptr) // [T, 1, B] for scoring (or nullptr) +{ + const auto& hp = model.hparams; + const auto& tensors = model.tensors; + const int D = hp.neck_dim; // 256 + const int NQ = hp.ddec_num_queries; // 200 + const int B = (int)enc_feats->ne[2]; // batch (1) + const int feat_hw = hp.n_img_embd(); // 72 - // ggml_flash_attn_ext expects mask storage in fp16. - return ggml_cont(ctx, ggml_cast(ctx, full_bias, GGML_TYPE_F16)); -} + // ── Initialize queries from query_embed ────────────────────────────── + auto* content = ggml_reshape_3d(ctx, model.ddec.query_embed, D, NQ, 1); + auto* pres_tok = ggml_reshape_3d(ctx, model.ddec.presence_token, D, 1, 1); + auto* queries = ggml_concat(ctx, pres_tok, content, 1); // [D, NQ+1, B=1] -/***************************************************************************** -** Geometry / exemplar encoder — graph building -*****************************************************************************/ + // Reference points: sigmoid → initial anchor boxes + auto* ref_pts_raw = ggml_cont(ctx, tensors.at("ddec.reference_points.weight")); // [4, NQ] + auto* ref_boxes = ggml_sigmoid(ctx, ref_pts_raw); // [4, NQ] + ref_boxes = ggml_reshape_3d(ctx, ref_boxes, 4, NQ, 1); // [4, NQ, 1] +#ifndef NDEBUG + auto* ref_boxes_dbg = ggml_cont(ctx, ref_boxes); + ggml_set_name(ref_boxes_dbg, "ddec_ref_boxes_init"); +#endif -// Sinusoidal positional encoding for box coordinates. -// Matches Python PositionEmbeddingSine.encode_boxes(cx, cy, w, h). -// Output: [258] = [pos_y(128), pos_x(128), h, w] -static void sam3_sine_encode_box(float* out, float cx, float cy, float w, float h, - int num_pos_feats, int temperature) { - const float scale = 2.0f * (float)M_PI; - float x_embed = cx * scale; - float y_embed = cy * scale; + // ── Run decoder layers ─────────────────────────────────────────────── + // Per-layer: recompute query_pos from updated ref_boxes (matching Python exactly) + struct ggml_tensor* last_presence = pres_tok; + for (int i = 0; i < hp.ddec_layers; ++i) { + // Recompute query_pos from current ref_boxes via sine embed + ref_point_head MLP + auto* query_pos = sam3_build_query_pos(ctx, model, ref_boxes, sine_dim_t, i); - // dim_t[i] = temperature^(2*(i//2)/num_pos_feats) - for (int i = 0; i < num_pos_feats; ++i) { - int div_idx = 2 * (i / 2); - float dim_t = powf((float)temperature, (float)div_idx / (float)num_pos_feats); - float px = x_embed / dim_t; - float py = y_embed / dim_t; - // Interleaved sin/cos: even indices get sin, odd get cos - if (i % 2 == 0) { - out[i] = sinf(py); // pos_y first - out[num_pos_feats + i] = sinf(px); // pos_x second - } else { - out[i] = cosf(py); - out[num_pos_feats + i] = cosf(px); + // Compute box-relative positional bias for image cross-attention + struct ggml_tensor* rpb_mask = nullptr; + if (rpb_coords) { + rpb_mask = sam3_compute_box_rpb(ctx, model, ref_boxes, rpb_coords, feat_hw, i); } - } - out[2 * num_pos_feats] = h; // h - out[2 * num_pos_feats + 1] = w; // w -} -// CPU-side ROI Align matching torchvision.ops.roi_align behavior. -// Features in ggml [C, W, H] layout. Box in XYXY format scaled to feature grid coords. -// Uses sub-sampling matching torchvision's sampling_ratio=0 (auto). -// Output: [C * roi_size * roi_size] in ggml layout. -static void sam3_roi_align_single( - const float* feats, // [C, W, H] ggml layout (C innermost, then W, then H) - int C, int W_feat, int H_feat, - float x0, float y0, float x1, float y1, - int roi_size, - float* out) // [C, roi_size, roi_size] -{ - float roi_w = std::max(x1 - x0, 1e-6f); - float roi_h = std::max(y1 - y0, 1e-6f); - float bin_w = roi_w / (float)roi_size; - float bin_h = roi_h / (float)roi_size; + queries = sam3_ddec_layer_forward(ctx, model.ddec.layers[i], + queries, query_pos, + enc_feats, enc_pos, + text_feats, hp.ddec_heads, + text_attn_bias, + rpb_mask, + i); - // Match torchvision sampling_ratio=0: use ceil(bin_size) sub-samples - int sample_y_count = std::max(1, (int)ceilf(bin_h)); - int sample_x_count = std::max(1, (int)ceilf(bin_w)); - float inv_count = 1.0f / (float)(sample_y_count * sample_x_count); + // Box refinement after each layer (on object queries only, not presence token) + auto* obj_q = ggml_view_3d(ctx, queries, D, NQ, 1, + queries->nb[1], queries->nb[2], 1 * queries->nb[1]); + obj_q = ggml_cont(ctx, obj_q); + sam3_name_tensorf(obj_q, "ddec_layer%d_out", i); - auto bilinear_sample = [&](float sx, float sy, int c) -> float { - // Clamp to [-0.5, W-0.5] then adjust — matching torchvision - if (sx < -1.0f || sx > (float)W_feat || sy < -1.0f || sy > (float)H_feat) - return 0.0f; + auto* pres_q = ggml_view_3d(ctx, queries, D, 1, 1, + queries->nb[1], queries->nb[2], 0); + pres_q = ggml_cont(ctx, pres_q); + if (i == 0) { + ggml_set_name(pres_q, "ddec_layer0_presence"); + } + last_presence = pres_q; - sy = std::max(0.0f, sy); - sx = std::max(0.0f, sx); + // Apply the final decoder norm before box refinement (use_normed_output_consistently) + auto* obj_q_normed = sam3_layer_norm(ctx, obj_q, + tensors.at("ddec.norm.weight"), + tensors.at("ddec.norm.bias")); - int y_lo = (int)sy; - int x_lo = (int)sx; - int y_hi = std::min(y_lo + 1, H_feat - 1); - int x_hi = std::min(x_lo + 1, W_feat - 1); - y_lo = std::min(y_lo, H_feat - 1); - x_lo = std::min(x_lo, W_feat - 1); + // Shared bbox_embed MLP + auto* bd = obj_q_normed; + for (int j = 0; j < 3; ++j) { + auto wn = "ddec.bbox_embed.layers." + std::to_string(j) + ".weight"; + auto bn = "ddec.bbox_embed.layers." + std::to_string(j) + ".bias"; + bd = ggml_mul_mat(ctx, tensors.at(wn), bd); + bd = ggml_add(ctx, bd, tensors.at(bn)); + if (j < 2) bd = ggml_relu(ctx, bd); + } + // bd: [4, NQ, 1] - float ly = sy - (float)y_lo; - float lx = sx - (float)x_lo; + // ref_boxes = sigmoid(inverse_sigmoid(ref_boxes) + box_delta) + auto* ref_inv_cur = sam3_inverse_sigmoid(ctx, ref_boxes); + ref_boxes = ggml_sigmoid(ctx, ggml_add(ctx, ref_inv_cur, bd)); + sam3_name_tensorf(ref_boxes, "ddec_layer%d_refboxes", i); + } - // ggml layout: feats[c + x * C + y * C * W_feat] - float v00 = feats[c + x_lo * C + y_lo * C * W_feat]; - float v10 = feats[c + x_hi * C + y_lo * C * W_feat]; - float v01 = feats[c + x_lo * C + y_hi * C * W_feat]; - float v11 = feats[c + x_hi * C + y_hi * C * W_feat]; + // ── Final normalization ────────────────────────────────────────────── + // Match Python: decoder.norm is applied to object queries only. + auto* obj_queries = ggml_view_3d(ctx, queries, D, NQ, 1, + queries->nb[1], queries->nb[2], 1 * queries->nb[1]); + obj_queries = ggml_cont(ctx, obj_queries); + obj_queries = sam3_layer_norm(ctx, obj_queries, + tensors.at("ddec.norm.weight"), + tensors.at("ddec.norm.bias")); + ggml_set_name(obj_queries, "ddec_normed_output"); - return v00 * (1 - ly) * (1 - lx) + v10 * (1 - ly) * lx + v01 * ly * (1 - lx) + v11 * ly * lx; - }; + auto* queries_for_seg = ggml_concat(ctx, last_presence, obj_queries, 1); - for (int ph = 0; ph < roi_size; ++ph) { - for (int pw = 0; pw < roi_size; ++pw) { - for (int c = 0; c < C; ++c) { - float sum = 0.0f; - for (int iy = 0; iy < sample_y_count; ++iy) { - float sy = y0 + bin_h * ((float)ph + ((float)iy + 0.5f) / (float)sample_y_count); - for (int ix = 0; ix < sample_x_count; ++ix) { - float sx = x0 + bin_w * ((float)pw + ((float)ix + 0.5f) / (float)sample_x_count); - sum += bilinear_sample(sx, sy, c); - } - } - out[c + pw * C + ph * C * roi_size] = sum * inv_count; - } - } + auto* class_scores = sam3_dot_product_scoring(ctx, model, obj_queries, text_feats, text_valid_mask); + // class_scores: [NQ, B] + + // ── Presence score ─────────────────────────────────────────────────── + // Presence token head: LN + 3-layer MLP (D→D→D→1) + auto* pres_out = sam3_layer_norm(ctx, last_presence, + tensors.at("ddec.presence_token_out_norm.weight"), + tensors.at("ddec.presence_token_out_norm.bias")); + + for (int j = 0; j < 3; ++j) { + auto wn = "ddec.presence_token_head.layers." + std::to_string(j) + ".weight"; + auto bn = "ddec.presence_token_head.layers." + std::to_string(j) + ".bias"; + pres_out = ggml_mul_mat(ctx, tensors.at(wn), pres_out); + pres_out = ggml_add(ctx, pres_out, tensors.at(bn)); + if (j < 2) pres_out = ggml_relu(ctx, pres_out); } + // Keep presence as raw logit (no sigmoid yet — applied during post-processing) + auto* presence_score = ggml_reshape_2d(ctx, pres_out, 1, 1); + // presence_score: [1, B] — raw logit + + sam3_ddec_output out; + out.queries = queries_for_seg; // [D, NQ+1, B] + out.presence_feats = last_presence; // [D, 1, B] + out.pred_boxes = ref_boxes; // [4, NQ, B] + out.class_scores = class_scores; // [NQ, B] + out.presence_score = presence_score; // [1, B] + + return out; } -// Build geometry encoder graph and pre-compute box embeddings on CPU. -// Returns the geometry features as a pre-computed input tensor [D, N_geo, 1] -// where N_geo = n_exemplar_boxes + 1 (CLS). -// For dummy prompts (no boxes), N_geo = 1 (just CLS token). -struct sam3_geom_result { - struct ggml_tensor* geo_feats; // [D, N_geo, 1] - int n_tokens; // N_geo -}; +/***************************************************************************** +** Segmentation head (MaskFormer) — graph building +*****************************************************************************/ -static sam3_geom_result sam3_build_geom_enc_graph( +// Build pixel decoder: progressively upsample FPN features. +// fpn_feats[0]: [D, 288, 288, B] (highest res) +// fpn_feats[1]: [D, 144, 144, B] +// fpn_feats[2]: [D, 72, 72, B] (lowest res) +// Returns: [D, 288, 288, B] pixel features +// +// Python PixelDecoder.forward: +// prev_fpn = backbone_feats[-1] (lowest res) +// for bb_feat in backbone_feats[:-1][::-1]: (iterate from second-lowest to highest) +// prev_fpn = bb_feat + F.interpolate(prev_fpn, size=bb_feat.shape[-2:], mode="nearest") +// prev_fpn = conv_layers[i](prev_fpn) # conv on the MERGED result +// prev_fpn = F.relu(norms[i](prev_fpn)) # GroupNorm then ReLU +// +// Python uses GroupNorm(8, 256) — we use ggml_group_norm which normalizes ne[2] +// (the channel dim) in groups. The conv output is [W, H, D, B] with D in ne[2]. +static struct ggml_tensor* sam3_pixel_decoder( struct ggml_context* ctx, const sam3_model& model, - const sam3_pcs_params& params, - struct ggml_tensor* img_feats, // [D, N_img, 1] where N_img = H*H - struct ggml_tensor* img_pe) // [D, N_img, 1] sinusoidal PE + struct ggml_tensor* fpn_feats[3]) // [D, W, H, B] at 3 scales { - const auto& ge = model.geom_enc; - const int D = model.hparams.neck_dim; // 256 - const int n_heads = 8; - const int n_boxes = (int)(params.pos_exemplars.size() + params.neg_exemplars.size()); - const int N_geo = n_boxes + 1; // +1 for CLS - - // Input tensor: pre-computed geometry embeddings after final_proj + norm [D, N_geo, 1] - // The caller pre-computes: final_proj(box_embeds + CLS) → LayerNorm - auto* x = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, N_geo, 1); - ggml_set_name(x, "geom_post_final_proj"); // also used as upload target - ggml_set_input(x); - - // Transformer layers (3 layers: self-attn + cross-attn + FFN, pre-norm) - for (int i = 0; i < (int)ge.layers.size(); ++i) { - const auto& ly = ge.layers[i]; - - // 1. Self-attention (pre-norm, pos_enc_at_attn=False) - { - auto* shortcut = x; - auto* xn = sam3_layer_norm(ctx, x, ly.norm1_w, ly.norm1_b); - - // Q = K = V = norm(x) — no positional encoding at self-attention - auto* Q = ggml_add(ctx, ggml_mul_mat(ctx, ggml_view_2d(ctx, ly.sa_in_proj_w, D, D, ly.sa_in_proj_w->nb[1], 0), xn), - ggml_view_1d(ctx, ly.sa_in_proj_b, D, 0)); - auto* K = ggml_add(ctx, ggml_mul_mat(ctx, ggml_view_2d(ctx, ly.sa_in_proj_w, D, D, ly.sa_in_proj_w->nb[1], D * ly.sa_in_proj_w->nb[1]), xn), - ggml_view_1d(ctx, ly.sa_in_proj_b, D, D * sizeof(float))); - auto* V = ggml_add(ctx, ggml_mul_mat(ctx, ggml_view_2d(ctx, ly.sa_in_proj_w, D, D, ly.sa_in_proj_w->nb[1], 2 * D * ly.sa_in_proj_w->nb[1]), xn), - ggml_view_1d(ctx, ly.sa_in_proj_b, D, 2 * D * sizeof(float))); - - const int64_t S = N_geo; - const int64_t HD = D / n_heads; + const auto& seg = model.seg_head; - Q = ggml_reshape_4d(ctx, Q, HD, n_heads, S, 1); - Q = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 1, 3)); - K = ggml_reshape_4d(ctx, K, HD, n_heads, S, 1); - K = ggml_cont(ctx, ggml_permute(ctx, K, 0, 2, 1, 3)); - V = ggml_reshape_4d(ctx, V, HD, n_heads, S, 1); - V = ggml_permute(ctx, V, 0, 2, 1, 3); + // Start from lowest resolution + auto* feat = fpn_feats[2]; // [D, 72, 72, B] - float scale = 1.0f / sqrtf((float)HD); - auto* sa_out = ggml_flash_attn_ext(ctx, Q, K, V, nullptr, scale, 0.0f, 0.0f); - sa_out = ggml_reshape_3d(ctx, sa_out, D, S, 1); + // Iteration 0: merge with FPN[1] (144x144) + // prev_fpn = FPN[1] + upsample(prev_fpn) + // Permute to [W, H, D, B] for conv operations + auto* prev = ggml_cont(ctx, ggml_permute(ctx, feat, 2, 0, 1, 3)); // [72, 72, D, B] + prev = ggml_upscale(ctx, prev, 2, GGML_SCALE_MODE_NEAREST); // [144, 144, D, B] + auto* fpn1 = ggml_cont(ctx, ggml_permute(ctx, fpn_feats[1], 2, 0, 1, 3)); // [144, 144, D, B] + prev = ggml_add(ctx, fpn1, prev); // merged + // Conv 3x3 on the MERGED result (not individual FPN feat) + prev = ggml_conv_2d_s1_ph(ctx, seg.up_conv_w[0], prev); + { + auto* b3d = ggml_reshape_3d(ctx, seg.up_conv_b[0], 1, 1, seg.up_conv_b[0]->ne[0]); + prev = ggml_add(ctx, prev, ggml_repeat(ctx, b3d, prev)); + } + // GroupNorm(8, 256) then ReLU — prev is [W, H, D, B] with D in ne[2] + prev = ggml_group_norm(ctx, prev, 8, 1e-5f); + { + auto* w3d = ggml_reshape_3d(ctx, seg.up_norm_w[0], 1, 1, seg.up_norm_w[0]->ne[0]); + prev = ggml_mul(ctx, prev, ggml_repeat(ctx, w3d, prev)); + auto* bn3d = ggml_reshape_3d(ctx, seg.up_norm_b[0], 1, 1, seg.up_norm_b[0]->ne[0]); + prev = ggml_add(ctx, prev, ggml_repeat(ctx, bn3d, prev)); + } + prev = ggml_relu(ctx, prev); + ggml_set_name(prev, "seg_pixel_dec_stage0"); - sa_out = ggml_mul_mat(ctx, ly.sa_out_proj_w, sa_out); - sa_out = ggml_add(ctx, sa_out, ly.sa_out_proj_b); + // Iteration 1: merge with FPN[0] (288x288) + prev = ggml_upscale(ctx, prev, 2, GGML_SCALE_MODE_NEAREST); // [288, 288, D, B] + auto* fpn0 = ggml_cont(ctx, ggml_permute(ctx, fpn_feats[0], 2, 0, 1, 3)); // [288, 288, D, B] + prev = ggml_add(ctx, fpn0, prev); // merged + // Conv 3x3 on the MERGED result + prev = ggml_conv_2d_s1_ph(ctx, seg.up_conv_w[1], prev); + { + auto* b3d = ggml_reshape_3d(ctx, seg.up_conv_b[1], 1, 1, seg.up_conv_b[1]->ne[0]); + prev = ggml_add(ctx, prev, ggml_repeat(ctx, b3d, prev)); + } + // GroupNorm(8, 256) then ReLU + prev = ggml_group_norm(ctx, prev, 8, 1e-5f); + { + auto* w3d = ggml_reshape_3d(ctx, seg.up_norm_w[1], 1, 1, seg.up_norm_w[1]->ne[0]); + prev = ggml_mul(ctx, prev, ggml_repeat(ctx, w3d, prev)); + auto* bn3d = ggml_reshape_3d(ctx, seg.up_norm_b[1], 1, 1, seg.up_norm_b[1]->ne[0]); + prev = ggml_add(ctx, prev, ggml_repeat(ctx, bn3d, prev)); + } + prev = ggml_relu(ctx, prev); + ggml_set_name(prev, "seg_pixel_dec_stage1"); - x = ggml_add(ctx, shortcut, sa_out); - } + // Python PixelDecoder allocates 3 conv layers but only uses 2 (one per + // upsample step). The 3rd conv (up_conv_w[2]) is unused. - // 2. Cross-attention (pre-norm, Q from x, K from img+PE, V from img) - { - auto* shortcut = x; - auto* xn = sam3_layer_norm(ctx, x, ly.norm2_w, ly.norm2_b); + auto* out = ggml_cont(ctx, ggml_permute(ctx, prev, 1, 2, 0, 3)); // [D, 288, 288, B] + return out; +} - // Q from normalized geometry tokens (no pos) - // K from image features + PE - // V from image features - auto* q_w = ggml_view_2d(ctx, ly.ca_q_w, D, D, ly.ca_q_w->nb[1], 0); - auto* k_w = ggml_view_2d(ctx, ly.ca_q_w, D, D, ly.ca_q_w->nb[1], D * ly.ca_q_w->nb[1]); - auto* v_w = ggml_view_2d(ctx, ly.ca_q_w, D, D, ly.ca_q_w->nb[1], 2 * D * ly.ca_q_w->nb[1]); - auto* q_b = ggml_view_1d(ctx, ly.ca_q_b, D, 0); - auto* k_b = ggml_view_1d(ctx, ly.ca_q_b, D, D * sizeof(float)); - auto* v_b = ggml_view_1d(ctx, ly.ca_q_b, D, 2 * D * sizeof(float)); +// Build the full segmentation head graph. +// +// Python UniversalSegmentationHead.forward: +// 1. Cross-attend encoder_hidden_states to prompt → updated encoder +// 2. _embed_pixels: replace lowest-res FPN feat with spatial portion of encoder output +// 3. Run pixel decoder on modified FPN feats +// 4. instance_seg_head (Conv1x1) +// 5. mask_predictor: einsum(mask_embed(queries), instance_embeds) +// +// enc_hidden: [D, N_spatial, B] — fusion encoder output (cross-attended in step 1) +// fpn_feats[3]: the 3 FPN features at different resolutions +// query_outputs: [D, N, B] selected object query outputs +// text_features: [D, T, B] for cross-attention (prompt) +// Returns: mask_logits [W*H, N, B] (raw logits, not sigmoid) +static struct ggml_tensor* sam3_build_seg_head_graph( + struct ggml_context* ctx, + const sam3_model& model, + struct ggml_tensor* enc_hidden, // [D, N_spatial, B] fusion encoder output + struct ggml_tensor* fpn_feats[3], // FPN features at 3 scales + struct ggml_tensor* query_outputs, // [D, N, B] + struct ggml_tensor* text_features, // [D, T, B] (for cross-attn, can be nullptr) + struct ggml_tensor* text_attn_bias = nullptr) { + const auto& seg = model.seg_head; + const auto& tensors = model.tensors; + const int64_t D = enc_hidden->ne[0]; // 256 + const int64_t B = enc_hidden->ne[2]; // 1 + const int64_t N = query_outputs->ne[1]; // number of selected queries - auto* k_input = ggml_add(ctx, img_feats, img_pe); // pos_enc_at_cross_attn_keys + auto* enc = enc_hidden; + if (text_features) { + auto* ca_norm = sam3_layer_norm(ctx, enc, + tensors.at("seg.cross_attn_norm.weight"), + tensors.at("seg.cross_attn_norm.bias")); - auto* Q = ggml_add(ctx, ggml_mul_mat(ctx, q_w, xn), q_b); - auto* K = ggml_add(ctx, ggml_mul_mat(ctx, k_w, k_input), k_b); - auto* V = ggml_add(ctx, ggml_mul_mat(ctx, v_w, img_feats), v_b); + auto* ca_mask = sam3_expand_token_attn_bias(ctx, text_attn_bias, enc->ne[1], 8, B); + auto* ca_out = sam3_multihead_attn_fused(ctx, ca_norm, text_features, + seg.ca_prompt_q_w, seg.ca_prompt_q_b, + seg.ca_prompt_out_w, seg.ca_prompt_out_b, + 8, ca_mask); + enc = ggml_add(ctx, enc, ca_out); + } + // enc: [D, N_spatial, B] + ggml_set_name(enc, "seg_enc_after_ca"); - const int64_t S_q = N_geo; - const int64_t S_kv = img_feats->ne[1]; - const int64_t HD = D / n_heads; + // Replace lowest-res FPN feat with spatial portion of encoder output + const int64_t feat_hw = model.hparams.n_img_embd(); // 72 + auto* enc_spatial = ggml_reshape_4d(ctx, enc, D, feat_hw, feat_hw, B); +#ifndef NDEBUG + auto* enc_spatial_dbg = ggml_cont(ctx, ggml_permute(ctx, enc_spatial, 2, 0, 1, 3)); + ggml_set_name(enc_spatial_dbg, "seg_enc_visual"); +#endif - Q = ggml_reshape_4d(ctx, Q, HD, n_heads, S_q, 1); - Q = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 1, 3)); - K = ggml_reshape_4d(ctx, K, HD, n_heads, S_kv, 1); - K = ggml_cont(ctx, ggml_permute(ctx, K, 0, 2, 1, 3)); - V = ggml_reshape_4d(ctx, V, HD, n_heads, S_kv, 1); - V = ggml_permute(ctx, V, 0, 2, 1, 3); + struct ggml_tensor* modified_fpn[3] = { + fpn_feats[0], + fpn_feats[1], + enc_spatial, // replaces original lowest-res FPN + }; - float scale = 1.0f / sqrtf((float)HD); - auto* ca_out = ggml_flash_attn_ext(ctx, Q, K, V, nullptr, scale, 0.0f, 0.0f); - ca_out = ggml_reshape_3d(ctx, ca_out, D, S_q, 1); + auto* pixel_feats = sam3_pixel_decoder(ctx, model, modified_fpn); +#ifndef NDEBUG + auto* pixel_feats_dbg = ggml_cont(ctx, ggml_permute(ctx, pixel_feats, 2, 0, 1, 3)); + ggml_set_name(pixel_feats_dbg, "seg_pixel_decoder_out"); +#endif - ca_out = ggml_mul_mat(ctx, ly.ca_out_w, ca_out); - ca_out = ggml_add(ctx, ca_out, ly.ca_out_b); + const int64_t W = pixel_feats->ne[1]; // 288 + const int64_t H = pixel_feats->ne[2]; // 288 - x = ggml_add(ctx, shortcut, ca_out); - } + // Instance segmentation head (Conv1x1) + auto* pf_conv = ggml_cont(ctx, ggml_permute(ctx, pixel_feats, 2, 0, 1, 3)); + pf_conv = ggml_conv_2d_sk_p0(ctx, tensors.at("seg.instance_seg_head.weight"), pf_conv); + { + auto* b3d = ggml_reshape_3d(ctx, tensors.at("seg.instance_seg_head.bias"), + 1, 1, tensors.at("seg.instance_seg_head.bias")->ne[0]); + pf_conv = ggml_add(ctx, pf_conv, ggml_repeat(ctx, b3d, pf_conv)); + } + auto* pixel_embed = ggml_cont(ctx, ggml_permute(ctx, pf_conv, 1, 2, 0, 3)); // [D, W, H, B] +#ifndef NDEBUG + auto* pixel_embed_dbg = ggml_cont(ctx, ggml_permute(ctx, pixel_embed, 2, 0, 1, 3)); + ggml_set_name(pixel_embed_dbg, "seg_instance_embed"); +#endif - // 3. FFN (pre-norm, ReLU) - { - auto* shortcut = x; - auto* xn = sam3_layer_norm(ctx, x, ly.norm3_w, ly.norm3_b); - auto* ffn = ggml_mul_mat(ctx, ly.ffn_fc1_w, xn); - ffn = ggml_add(ctx, ffn, ly.ffn_fc1_b); - ffn = ggml_relu(ctx, ffn); - ffn = ggml_mul_mat(ctx, ly.ffn_fc2_w, ffn); - ffn = ggml_add(ctx, ffn, ly.ffn_fc2_b); - x = ggml_add(ctx, shortcut, ffn); - } - sam3_name_tensorf(x, "geom_layer%d_out", i); + // Mask embedding MLP + auto* mask_embed = query_outputs; + for (int j = 0; j < 3; ++j) { + auto wn = "seg.mask_predictor.mask_embed.layers." + std::to_string(j) + ".weight"; + auto bn = "seg.mask_predictor.mask_embed.layers." + std::to_string(j) + ".bias"; + mask_embed = ggml_mul_mat(ctx, tensors.at(wn), mask_embed); + mask_embed = ggml_add(ctx, mask_embed, tensors.at(bn)); + if (j < 2) mask_embed = ggml_relu(ctx, mask_embed); } + ggml_set_name(mask_embed, "seg_mask_embed"); - // Final encode norm - x = sam3_layer_norm(ctx, x, ge.encode_norm_w, ge.encode_norm_b); - ggml_set_name(x, "geom_output"); - ggml_set_output(x); + // Mask prediction: einsum('bqc,bchw->bqhw') + auto* pe_flat = ggml_reshape_3d(ctx, pixel_embed, D, W * H, B); + auto* masks = ggml_mul_mat(ctx, pe_flat, mask_embed); + ggml_set_name(masks, "seg_mask_logits"); - sam3_geom_result result; - result.geo_feats = x; - result.n_tokens = N_geo; - return result; + return masks; } -// Pre-compute geometry encoder input on CPU: box embeddings + CLS token. -// Reads model weights from GPU, computes embeddings, returns float vector. -// Layout: [D, N_geo] row-major where N_geo = n_boxes + 1 (CLS last). -static std::vector sam3_precompute_geom_input( - const sam3_model& model, - const sam3_pcs_params& params, - const float* img_feats_data, // [D, W, H] ggml-layout backbone features (nullable if no boxes) - int W_feat, int H_feat) { - const auto& ge = model.geom_enc; - const int D = model.hparams.neck_dim; // 256 - const int roi_size = 7; - const int num_pos_feats = D / 2; // 128 +/***************************************************************************** +** Memory encoder (Phase 7, Step 7.1) +*****************************************************************************/ - // Collect all exemplar boxes - struct box_info { - float cx, cy, w, h; - int label; - }; - std::vector boxes; - for (const auto& b : params.pos_exemplars) { - // API provides XYXY in original image space — convert to normalized CxCyWH [0,1] - float cx = (b.x0 + b.x1) * 0.5f; - float cy = (b.y0 + b.y1) * 0.5f; - float bw = b.x1 - b.x0; - float bh = b.y1 - b.y0; - boxes.push_back({cx, cy, bw, bh, 0}); // label 0 = positive - } - for (const auto& b : params.neg_exemplars) { - float cx = (b.x0 + b.x1) * 0.5f; - float cy = (b.y0 + b.y1) * 0.5f; - float bw = b.x1 - b.x0; - float bh = b.y1 - b.y0; - boxes.push_back({cx, cy, bw, bh, 1}); // label 1 = negative - } +// CXBlock: depthwise conv + LayerNorm + pointwise MLP with residual scaling. +static struct ggml_tensor* sam3_cxblock_forward( + struct ggml_context* ctx, + struct ggml_tensor* x, // [D, H, W, B] + struct ggml_tensor* dw_w, // [7, 7, 1, D] depthwise + struct ggml_tensor* dw_b, // [D] + struct ggml_tensor* norm_w, // [D] + struct ggml_tensor* norm_b, // [D] + struct ggml_tensor* fc1_w, // [D, 1024] + struct ggml_tensor* fc1_b, // [1024] + struct ggml_tensor* fc2_w, // [1024, D] + struct ggml_tensor* fc2_b, // [D] + struct ggml_tensor* gamma) // [D] +{ + const int D = (int)x->ne[0]; + const int H = (int)x->ne[1]; + const int W = (int)x->ne[2]; - const int n_boxes = (int)boxes.size(); - const int N_geo = n_boxes + 1; + // ggml conv expects WHCB layout; internal feature maps are stored as CWHB. + auto* x_whcb = ggml_cont(ctx, ggml_permute(ctx, x, 2, 0, 1, 3)); - // Read needed weights from GPU - std::vector box_proj_w_data(4 * D), box_proj_b_data(D); - std::vector type_embed_data(D * 2); - std::vector cls_data(D); - std::vector box_pos_proj_w_data(258 * D), box_pos_proj_b_data(D); - std::vector box_pool_proj_w_data(7 * 7 * D * D), box_pool_proj_b_data(D); - std::vector img_pre_norm_w_data(D), img_pre_norm_b_data(D); + // Depthwise conv (groups = D): use the direct depthwise path. + // ggml_conv_2d_dw_direct only supports f32 kernel — cast if needed. + auto* dw_w_f32 = (dw_w->type == GGML_TYPE_F32) ? dw_w : ggml_cast(ctx, dw_w, GGML_TYPE_F32); + auto* h = ggml_conv_2d_dw_direct(ctx, dw_w_f32, x_whcb, 1, 1, 3, 3, 1, 1); + h = ggml_add(ctx, h, ggml_reshape_4d(ctx, dw_b, 1, 1, D, 1)); + h = ggml_cont(ctx, ggml_permute(ctx, h, 1, 2, 0, 3)); - auto read_f32 = [](struct ggml_tensor* t, float* dst, size_t n) { - sam3_read_f32(t, dst, n); - }; + // LayerNorm2d + h = sam3_layer_norm_2d(ctx, h, norm_w, norm_b); - read_f32(ge.box_proj_w, box_proj_w_data.data(), 4 * D); - read_f32(ge.box_proj_b, box_proj_b_data.data(), D); - read_f32(ge.type_embed, type_embed_data.data(), D * 2); - read_f32(ge.cls_token, cls_data.data(), D); - read_f32(ge.box_pos_proj_w, box_pos_proj_w_data.data(), 258 * D); - read_f32(ge.box_pos_proj_b, box_pos_proj_b_data.data(), D); + // Pointwise MLP: reshape to [D, H*W, B], apply FC, reshape back + auto* flat = ggml_reshape_3d(ctx, h, D, H * W, 1); + flat = ggml_add(ctx, ggml_mul_mat(ctx, fc1_w, flat), fc1_b); + flat = ggml_gelu(ctx, flat); + flat = ggml_add(ctx, ggml_mul_mat(ctx, fc2_w, flat), fc2_b); + h = ggml_reshape_4d(ctx, flat, D, H, W, 1); - if (n_boxes > 0) { - read_f32(ge.box_pool_proj_w, box_pool_proj_w_data.data(), 7 * 7 * D * D); - read_f32(ge.box_pool_proj_b, box_pool_proj_b_data.data(), D); - read_f32(ge.img_pre_norm_w, img_pre_norm_w_data.data(), D); - read_f32(ge.img_pre_norm_b, img_pre_norm_b_data.data(), D); - } + // Residual with learnable scaling: x + gamma * h + auto* gamma_4d = ggml_reshape_4d(ctx, gamma, D, 1, 1, 1); + h = ggml_mul(ctx, h, gamma_4d); + return ggml_add(ctx, x, h); +} - // Output: [D * N_geo] in row-major [D, N_geo] ggml order - std::vector out(D * N_geo, 0.0f); +/***************************************************************************** +** Prompt building for memory attention (Phase 7) +*****************************************************************************/ - // Encode each box - for (int bi = 0; bi < n_boxes; ++bi) { - const auto& box = boxes[bi]; - float embed[256] = {}; +struct sam3_prompt_data { + std::vector prompt; // [MD * M_total] flat + std::vector prompt_pos; // [MD * M_total] flat + int M_total; // total prompt tokens + int num_obj_ptr_tokens; // pointer tokens at end (excluded from RoPE) + int M_spatial; // spatial memory tokens = M_total - num_obj_ptr_tokens +}; - // 1. Direct projection: Linear(4, D) - // box_proj_w: ggml [4, D] = PyTorch [D, 4] - // out = x @ W^T + b where x=[cx,cy,w,h] - { - float coords[4] = {box.cx, box.cy, box.w, box.h}; - for (int d = 0; d < D; ++d) { - float sum = box_proj_b_data[d]; - for (int j = 0; j < 4; ++j) - sum += coords[j] * box_proj_w_data[j + d * 4]; // ggml: [4, D] stride - embed[d] = sum; - } - } +// Build prompt and prompt_pos tensors CPU-side for memory attention. +// mem_slot_feats[i]: [MD * H * H] raw spatial features (ggml layout) +// mem_slot_pes[i]: [MD * H * H] sinusoidal spatial PE (ggml layout) +// spatial_tpos[i]: temporal position for slot i (0 = conditioning, >0 = non-cond t_pos) +// obj_ptrs[i]: [D] object pointer data +// ptr_tpos[i]: temporal position for pointer i (<0 = skip) +static sam3_prompt_data sam3_build_prompt_and_pos( + const sam3_model& model, + const std::vector>& mem_slot_feats, + const std::vector>& mem_slot_pes, + const std::vector& spatial_tpos, + const std::vector>& obj_ptrs, + const std::vector& ptr_tpos, + int eff_feat_size = 0) { + const auto& hp = model.hparams; + const int MD = hp.mem_out_dim; // 64 + const int D = hp.neck_dim; // 256 + const int H = (eff_feat_size > 0) ? eff_feat_size : hp.feat_size(); + // Per-slot token count: from actual data size (handles both HxH and perceiver 512) + const int HH = mem_slot_feats.empty() ? (H * H) + : (int)(mem_slot_feats[0].size() / MD); - // 2. ROI Align + Conv2d(D, D, 7) pool projection - if (img_feats_data) { - // Apply img_pre_norm (LayerNorm) to image features before pooling - // For simplicity, we compute LayerNorm per spatial position - // Actually, LayerNorm is applied to the [D] dimension at each spatial location - // But we need the full feature map, so let's normalize in-place copy - const int N_spatial = W_feat * H_feat; - std::vector normed_feats(D * N_spatial); + sam3_prompt_data pd; + pd.num_obj_ptr_tokens = 0; - for (int s = 0; s < N_spatial; ++s) { - // Compute mean and variance over D dimension - float mean = 0.0f; - for (int d = 0; d < D; ++d) - mean += img_feats_data[d + s * D]; - mean /= D; + // Read maskmem_tpos_enc from model (one tensor [MD, 1, 1, 7]) + std::vector tpos_all(MD * hp.num_maskmem); + if (model.mem_enc.tpos[0]) { + sam3_read_f32(model.mem_enc.tpos[0], tpos_all.data(), MD * hp.num_maskmem); + } - float var = 0.0f; - for (int d = 0; d < D; ++d) { - float diff = img_feats_data[d + s * D] - mean; - var += diff * diff; - } - var /= D; + // 1. Spatial memory tokens + for (size_t s = 0; s < mem_slot_feats.size(); ++s) { + const int slot_tokens = (int)(mem_slot_feats[s].size() / MD); + // Append spatial features + pd.prompt.insert(pd.prompt.end(), mem_slot_feats[s].begin(), mem_slot_feats[s].end()); - float inv_std = 1.0f / sqrtf(var + 1e-5f); - for (int d = 0; d < D; ++d) { - normed_feats[d + s * D] = - (img_feats_data[d + s * D] - mean) * inv_std * img_pre_norm_w_data[d] + img_pre_norm_b_data[d]; - } - } + // Build positional encoding: spatial PE + temporal PE + std::vector pos = mem_slot_pes[s]; // copy spatial PE + int tpos_idx = spatial_tpos[s]; + if (tpos_idx >= 0 && tpos_idx < hp.num_maskmem) { + int enc_idx = hp.num_maskmem - tpos_idx - 1; // Python: maskmem_tpos_enc[num_maskmem - tpos - 1] + for (int n = 0; n < slot_tokens; ++n) + for (int d = 0; d < MD; ++d) + pos[d + n * MD] += tpos_all[d + enc_idx * MD]; + } + pd.prompt_pos.insert(pd.prompt_pos.end(), pos.begin(), pos.end()); + } + pd.M_spatial = (int)mem_slot_feats.size() * HH; - // Convert CxCyWH [0,1] → XYXY in feature grid coordinates - float fx0 = (box.cx - box.w * 0.5f) * (float)W_feat; - float fy0 = (box.cy - box.h * 0.5f) * (float)H_feat; - float fx1 = (box.cx + box.w * 0.5f) * (float)W_feat; - float fy1 = (box.cy + box.h * 0.5f) * (float)H_feat; + // 2. Object pointer tokens (split each D=256 pointer into D/MD=4 tokens of MD=64) + const int split = D / MD; // 4 - // ROI Align - std::vector roi_data(D * roi_size * roi_size); - sam3_roi_align_single(normed_feats.data(), D, W_feat, H_feat, - fx0, fy0, fx1, fy1, roi_size, roi_data.data()); + // Read obj_ptr_tpos_proj weights for CPU-side matmul + std::vector tpos_w(D * MD), tpos_b(MD); + if (model.obj_ptr_tpos_w) { + sam3_read_f32(model.obj_ptr_tpos_w, tpos_w.data(), D * MD); + sam3_read_f32(model.obj_ptr_tpos_b, tpos_b.data(), MD); + } - // Conv2d(D, D, 7): kernel [7, 7, D, D] in ggml = [D_out, D_in, kH, kW] in PyTorch - // Since roi_size = 7 = kernel_size, output is [D, 1, 1] - // This is effectively a matrix multiply: out[d_out] = sum over (d_in, kh, kw) - for (int d_out = 0; d_out < D; ++d_out) { - float sum = box_pool_proj_b_data[d_out]; - for (int kh = 0; kh < roi_size; ++kh) { - for (int kw = 0; kw < roi_size; ++kw) { - for (int d_in = 0; d_in < D; ++d_in) { - // ggml weight layout: [kW=7, kH=7, D_in=256, D_out=256] - int w_idx = kw + kh * 7 + d_in * 7 * 7 + d_out * 7 * 7 * D; - // roi_data layout: [C, W, H] = [D_in, kW, kH] - int r_idx = d_in + kw * D + kh * D * roi_size; - sum += box_pool_proj_w_data[w_idx] * roi_data[r_idx]; - } - } - } - embed[d_out] += sum; - } - } + for (size_t p = 0; p < obj_ptrs.size(); ++p) { + int rel = ptr_tpos[p]; + if (rel < 0) continue; - // 3. Sinusoidal positional encoding + Linear(258, D) - { - float pos_enc[258]; - sam3_sine_encode_box(pos_enc, box.cx, box.cy, box.w, box.h, - num_pos_feats, 10000); + // 1D sine PE for temporal position (normalized by max_obj_ptrs - 1 = 15) + std::vector sine_pe(D); + sam3_get_1d_sine_pe(sine_pe.data(), (float)rel / 15.0f, D); - for (int d = 0; d < D; ++d) { - float sum = box_pos_proj_b_data[d]; - for (int j = 0; j < 258; ++j) - sum += pos_enc[j] * box_pos_proj_w_data[j + d * 258]; - embed[d] += sum; - } + // Project 256-dim sine PE → 64-dim via obj_ptr_tpos_proj (CPU matmul) + // W is [D=256, MD=64] in ggml; y[j] = sum_i W[i + j*D] * x[i] + b[j] + std::vector proj_pe(MD, 0.0f); + for (int j = 0; j < MD; ++j) { + float sum = tpos_b[j]; + for (int i = 0; i < D; ++i) + sum += tpos_w[i + j * D] * sine_pe[i]; + proj_pe[j] = sum; } - // 4. Label embedding - { - int label = box.label; - for (int d = 0; d < D; ++d) - embed[d] += type_embed_data[d + label * D]; + // Split pointer [256] → 4 tokens of [64] + // Python: ptr.view(1,1,4,64).permute(1,2,0,3).reshape(4,1,64) + const auto& ptr = obj_ptrs[p]; + for (int t = 0; t < split; ++t) { + for (int d = 0; d < MD; ++d) { + pd.prompt.push_back(ptr[t * MD + d]); + pd.prompt_pos.push_back(proj_pe[d]); + } + pd.num_obj_ptr_tokens++; } + } - // Store in output: ggml layout [D, N_geo] — embed goes at column bi - for (int d = 0; d < D; ++d) - out[d + bi * D] = embed[d]; - } - - // CLS token goes at position n_boxes (last) - for (int d = 0; d < D; ++d) - out[d + n_boxes * D] = cls_data[d]; - - // Apply final_proj (Linear(D,D)) + LayerNorm on CPU - std::vector proj_w(D * D), proj_b(D); - std::vector ln_w(D), ln_b(D); - read_f32(ge.post_proj_w, proj_w.data(), D * D); - read_f32(ge.post_proj_b, proj_b.data(), D); - read_f32(ge.norm_w, ln_w.data(), D); - read_f32(ge.norm_b, ln_b.data(), D); - - std::vector projected(D * N_geo); - for (int t = 0; t < N_geo; ++t) { - // Linear: y = W @ x + b - // W is stored in PyTorch row-major [D_out, D_in] → ggml [D_in, D_out] - // proj_w[d_in + d_out * D] = W_py[d_out, d_in] - for (int d_out = 0; d_out < D; ++d_out) { - float sum = proj_b[d_out]; - for (int d_in = 0; d_in < D; ++d_in) - sum += out[d_in + t * D] * proj_w[d_in + d_out * D]; - projected[d_out + t * D] = sum; - } - } - if (n_boxes > 0) { - fprintf(stderr, "%s: %d exemplar boxes encoded (%d pos, %d neg)\n", - __func__, n_boxes, - (int)params.pos_exemplars.size(), - (int)params.neg_exemplars.size()); - } - - // LayerNorm over D dimension for each token - for (int t = 0; t < N_geo; ++t) { - float mean = 0.0f; - for (int d = 0; d < D; ++d) - mean += projected[d + t * D]; - mean /= D; - - float var = 0.0f; - for (int d = 0; d < D; ++d) { - float diff = projected[d + t * D] - mean; - var += diff * diff; - } - var /= D; - - float inv_std = 1.0f / sqrtf(var + 1e-5f); - for (int d = 0; d < D; ++d) { - float normalized = (projected[d + t * D] - mean) * inv_std; - out[d + t * D] = normalized * ln_w[d] + ln_b[d]; - } - } - return out; + pd.M_total = pd.M_spatial + pd.num_obj_ptr_tokens; + return pd; } /***************************************************************************** -** Fusion encoder — graph building (6 layers) +** Memory attention (Phase 7, Step 7.2) *****************************************************************************/ -// Single fusion encoder layer. -// x: [D, N, B] image features (N=5184), prompt: [D, T, B] text/exemplar tokens, pos: [D, N, B] -// Returns: updated x [D, N, B] -static struct ggml_tensor* sam3_fenc_layer_forward( +// Build memory attention graph. +// curr_tokens: [D, N, 1] — current frame image tokens (flattened from 72x72 = 5184) +// src_pos: [D, N, 1] — sinusoidal PE for current tokens (pos_enc_at_input) +// prompt: [MD, M_total, 1] — concatenated memory features + split pointer tokens +// prompt_pos: [MD, M_total, 1] — positional encodings for prompt +// rope_freqs: [2, D/2, N] — axial RoPE frequencies for Q and self-attn K +// rope_k_freqs:[2, D/2, M_spatial] — repeated axial RoPE for cross-attn spatial K +// num_obj_ptr_tokens: number of pointer tokens at end of prompt (excluded from RoPE) +// Returns: conditioned tokens [D, N, 1] +static struct ggml_tensor* sam3_build_mem_attn_graph( struct ggml_context* ctx, - const sam3_fenc_layer& ly, - struct ggml_tensor* x, - struct ggml_tensor* prompt, - struct ggml_tensor* pos, - struct ggml_tensor* prompt_attn_bias, - int n_heads) { - // Self-attention: Q/K get positional encoding, V does not - { - auto* shortcut = x; - auto* x_norm = sam3_layer_norm(ctx, x, ly.norm1_w, ly.norm1_b); - auto* q_in = ggml_add(ctx, x_norm, pos); - auto* k_in = ggml_add(ctx, x_norm, pos); - - const int64_t D = x->ne[0]; + const sam3_model& model, + struct ggml_tensor* curr_tokens, // [D, N, 1] + struct ggml_tensor* src_pos, // [D, N, 1] + struct ggml_tensor* prompt, // [MD, M_total, 1] + struct ggml_tensor* prompt_pos, // [MD, M_total, 1] + struct ggml_tensor* rope_freqs, // [2, D/2, N] + struct ggml_tensor* rope_k_freqs, // [2, D/2, M_spatial] or nullptr + int num_obj_ptr_tokens) { + const auto& ma = model.mem_attn; + const int D = model.hparams.neck_dim; // 256 + const int N = (int)curr_tokens->ne[1]; + const int M_total = (int)prompt->ne[1]; + const int M_spatial = M_total - num_obj_ptr_tokens; - auto* q_w = ggml_view_2d(ctx, ly.sa_in_proj_w, D, D, ly.sa_in_proj_w->nb[1], 0); - auto* k_w = ggml_view_2d(ctx, ly.sa_in_proj_w, D, D, ly.sa_in_proj_w->nb[1], D * ly.sa_in_proj_w->nb[1]); - auto* v_w = ggml_view_2d(ctx, ly.sa_in_proj_w, D, D, ly.sa_in_proj_w->nb[1], 2 * D * ly.sa_in_proj_w->nb[1]); + // pos_enc_at_input: x = curr + 0.1 * src_pos + auto* x = ggml_add(ctx, curr_tokens, ggml_scale(ctx, src_pos, 0.1f)); + ggml_set_name(x, "phase7_mem_attn_input"); - auto* q_b = ggml_view_1d(ctx, ly.sa_in_proj_b, D, 0); - auto* k_b = ggml_view_1d(ctx, ly.sa_in_proj_b, D, D * sizeof(float)); - auto* v_b = ggml_view_1d(ctx, ly.sa_in_proj_b, D, 2 * D * sizeof(float)); + for (int l = 0; l < (int)ma.layers.size(); ++l) { + const auto& ly = ma.layers[l]; - auto* Q = ggml_add(ctx, ggml_mul_mat(ctx, q_w, q_in), q_b); - auto* K = ggml_add(ctx, ggml_mul_mat(ctx, k_w, k_in), k_b); - auto* V = ggml_add(ctx, ggml_mul_mat(ctx, v_w, x_norm), v_b); + // ── Self-attention with RoPE ────────────────────────────────────── + { + auto* x_norm = sam3_layer_norm(ctx, x, ly.norm1_w, ly.norm1_b); + auto* q = ggml_add(ctx, ggml_mul_mat(ctx, ly.sa_q_w, x_norm), ly.sa_q_b); + auto* k = ggml_add(ctx, ggml_mul_mat(ctx, ly.sa_k_w, x_norm), ly.sa_k_b); + auto* v = ggml_add(ctx, ggml_mul_mat(ctx, ly.sa_v_w, x_norm), ly.sa_v_b); - const int64_t N = x->ne[1]; - const int64_t B = x->ne[2]; - const int64_t HD = D / n_heads; + // Apply RoPE to Q and K (single-head, [D, N, 1]) + q = sam3_apply_rope(ctx, q, rope_freqs); + k = sam3_apply_rope(ctx, k, rope_freqs); - Q = ggml_reshape_4d(ctx, Q, HD, n_heads, N, B); - Q = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 1, 3)); - K = ggml_reshape_4d(ctx, K, HD, n_heads, N, B); - K = ggml_cont(ctx, ggml_permute(ctx, K, 0, 2, 1, 3)); - V = ggml_reshape_4d(ctx, V, HD, n_heads, N, B); - V = ggml_permute(ctx, V, 0, 2, 1, 3); + // Single-head attention (256-dim) + q = ggml_reshape_4d(ctx, q, D, 1, N, 1); + k = ggml_reshape_4d(ctx, k, D, 1, N, 1); + v = ggml_reshape_4d(ctx, v, D, 1, N, 1); + q = ggml_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3)); + k = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3)); + v = ggml_permute(ctx, v, 0, 2, 1, 3); - float scale = 1.0f / sqrtf((float)HD); - auto* sa_out = ggml_flash_attn_ext(ctx, Q, K, V, nullptr, scale, 0.0f, 0.0f); - sa_out = ggml_reshape_3d(ctx, sa_out, D, N, B); + float scale = 1.0f / sqrtf((float)D); + auto* sa_out = ggml_flash_attn_ext(ctx, q, k, v, nullptr, scale, 0.0f, 0.0f); + sa_out = ggml_reshape_3d(ctx, sa_out, D, N, 1); + sa_out = ggml_add(ctx, ggml_mul_mat(ctx, ly.sa_out_w, sa_out), ly.sa_out_b); + x = ggml_add(ctx, x, sa_out); + sam3_name_tensorf(x, "phase7_mem_attn_layer%d_after_sa", l); + } - sa_out = ggml_mul_mat(ctx, ly.sa_out_proj_w, sa_out); - sa_out = ggml_add(ctx, sa_out, ly.sa_out_proj_b); + // ── Cross-attention to memory (kv_dim=64→256) with RoPE ────────── + { + auto* x_norm = sam3_layer_norm(ctx, x, ly.norm2_w, ly.norm2_b); + auto* q = ggml_add(ctx, ggml_mul_mat(ctx, ly.ca_q_w, x_norm), ly.ca_q_b); - x = ggml_add(ctx, shortcut, sa_out); - } + // K: project from (prompt + prompt_pos) — pos enc added before projection + auto* kv_with_pos = ggml_add(ctx, prompt, prompt_pos); + auto* k = ggml_add(ctx, ggml_mul_mat(ctx, ly.ca_k_w, kv_with_pos), ly.ca_k_b); + // V: project from prompt only (no pos enc) + auto* v = ggml_add(ctx, ggml_mul_mat(ctx, ly.ca_v_w, prompt), ly.ca_v_b); - // Cross-attention: Q from image features, K/V from prompt tokens. - // ca_q_w stores fused [D, 3*D] weights split as Q-proj, K-proj, V-proj. - { - auto* shortcut = x; - auto* x_norm = sam3_layer_norm(ctx, x, ly.norm2_w, ly.norm2_b); - const int64_t D = x->ne[0]; + // Apply RoPE to Q + q = sam3_apply_rope(ctx, q, rope_freqs); - auto* q_w = ggml_view_2d(ctx, ly.ca_q_w, D, D, ly.ca_q_w->nb[1], 0); - auto* k_w = ggml_view_2d(ctx, ly.ca_q_w, D, D, ly.ca_q_w->nb[1], D * ly.ca_q_w->nb[1]); - auto* v_w = ggml_view_2d(ctx, ly.ca_q_w, D, D, ly.ca_q_w->nb[1], 2 * D * ly.ca_q_w->nb[1]); + // Apply RoPE to spatial K only (exclude num_obj_ptr_tokens tail keys) + if (M_spatial > 0 && rope_k_freqs) { + auto* k_spatial = ggml_cont(ctx, ggml_view_3d(ctx, k, + D, M_spatial, 1, k->nb[1], k->nb[2], 0)); + k_spatial = sam3_apply_rope(ctx, k_spatial, rope_k_freqs); + if (num_obj_ptr_tokens > 0) { + auto* k_ptr = ggml_cont(ctx, ggml_view_3d(ctx, k, + D, num_obj_ptr_tokens, 1, k->nb[1], k->nb[2], + (size_t)M_spatial * k->nb[1])); + k = ggml_concat(ctx, k_spatial, k_ptr, 1); + } else { + k = k_spatial; + } + } - auto* q_b = ggml_view_1d(ctx, ly.ca_q_b, D, 0); - auto* k_b = ggml_view_1d(ctx, ly.ca_q_b, D, D * sizeof(float)); - auto* v_b = ggml_view_1d(ctx, ly.ca_q_b, D, 2 * D * sizeof(float)); + // Flash attention (single head) + q = ggml_reshape_4d(ctx, q, D, 1, N, 1); + k = ggml_reshape_4d(ctx, k, D, 1, M_total, 1); + v = ggml_reshape_4d(ctx, v, D, 1, M_total, 1); + q = ggml_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3)); + k = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3)); + v = ggml_permute(ctx, v, 0, 2, 1, 3); - auto* Q = ggml_add(ctx, ggml_mul_mat(ctx, q_w, x_norm), q_b); - auto* K = ggml_add(ctx, ggml_mul_mat(ctx, k_w, prompt), k_b); - auto* V = ggml_add(ctx, ggml_mul_mat(ctx, v_w, prompt), v_b); + float scale = 1.0f / sqrtf((float)D); + auto* ca_out = ggml_flash_attn_ext(ctx, q, k, v, nullptr, scale, 0.0f, 0.0f); + ca_out = ggml_reshape_3d(ctx, ca_out, D, N, 1); + ca_out = ggml_add(ctx, ggml_mul_mat(ctx, ly.ca_out_w, ca_out), ly.ca_out_b); + x = ggml_add(ctx, x, ca_out); + sam3_name_tensorf(x, "phase7_mem_attn_layer%d_after_ca", l); + } - const int64_t N_q = x->ne[1]; - const int64_t N_kv = prompt->ne[1]; - const int64_t B = x->ne[2]; - const int64_t HD = D / n_heads; + // ── FFN ─────────────────────────────────────────────────────────── + { + auto* x_norm = sam3_layer_norm(ctx, x, ly.norm3_w, ly.norm3_b); + auto* ffn = ggml_add(ctx, ggml_mul_mat(ctx, ly.ffn_fc1_w, x_norm), ly.ffn_fc1_b); + ffn = ggml_relu(ctx, ffn); + ffn = ggml_add(ctx, ggml_mul_mat(ctx, ly.ffn_fc2_w, ffn), ly.ffn_fc2_b); + x = ggml_add(ctx, x, ffn); + sam3_name_tensorf(x, "phase7_mem_attn_layer%d_after_ffn", l); + } + } - Q = ggml_reshape_4d(ctx, Q, HD, n_heads, N_q, B); - Q = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 1, 3)); - K = ggml_reshape_4d(ctx, K, HD, n_heads, N_kv, B); - K = ggml_cont(ctx, ggml_permute(ctx, K, 0, 2, 1, 3)); - V = ggml_reshape_4d(ctx, V, HD, n_heads, N_kv, B); - V = ggml_permute(ctx, V, 0, 2, 1, 3); + // Final norm + auto* norm_w = model.tensors.at("mem_attn.norm.weight"); + auto* norm_b = model.tensors.at("mem_attn.norm.bias"); + x = sam3_layer_norm(ctx, x, norm_w, norm_b); + ggml_set_name(x, "phase7_mem_attn_output"); - auto* ca_mask = sam3_expand_token_attn_bias(ctx, prompt_attn_bias, N_q, n_heads, B); - float scale = 1.0f / sqrtf((float)HD); - auto* ca_out = ggml_flash_attn_ext(ctx, Q, K, V, ca_mask, scale, 0.0f, 0.0f); - ca_out = ggml_reshape_3d(ctx, ca_out, D, N_q, B); + return x; // [D, N, 1] +} - ca_out = ggml_mul_mat(ctx, ly.ca_out_w, ca_out); - ca_out = ggml_add(ctx, ca_out, ly.ca_out_b); +/***************************************************************************** +** Object pointer extraction (Phase 7, Step 7.3) +*****************************************************************************/ - x = ggml_add(ctx, shortcut, ca_out); - } +// Extract object pointer from SAM output token via 3-layer MLP (CPU-side). +static void sam3_extract_obj_ptr_cpu( + const sam3_model& model, + const float* sam_token_data, // [D] + float obj_score, + float* out_ptr) // [D] +{ + const auto& hp = model.hparams; + const int D = hp.neck_dim; - { - auto* shortcut = x; - auto* x_norm = sam3_layer_norm(ctx, x, ly.norm3_w, ly.norm3_b); + // SAM2/EdgeTAM with fixed_no_obj_ptr: blend projected ptr with no_obj_ptr + // based on presence score λ. + // SAM3 / SAM2 without fixed_no_obj_ptr: binary threshold. + if ((hp.is_sam2() || hp.is_edgetam()) && hp.fixed_no_obj_ptr) { + // λ = (obj_score > 0) ? 1.0 : 0.0 (hard threshold, not sigmoid) + float lambda = (obj_score > 0.0f) ? 1.0f : 0.0f; - auto* ffn = ggml_mul_mat(ctx, ly.ffn_fc1_w, x_norm); - ffn = ggml_add(ctx, ffn, ly.ffn_fc1_b); - ffn = ggml_relu(ctx, ffn); - ffn = ggml_mul_mat(ctx, ly.ffn_fc2_w, ffn); - ffn = ggml_add(ctx, ffn, ly.ffn_fc2_b); + // Project token through MLP + std::vector h(D), tmp(D); + std::copy(sam_token_data, sam_token_data + D, h.data()); + for (int j = 0; j < 3; ++j) { + auto* w = model.obj_ptr_proj_w[j]; + auto* b = model.obj_ptr_proj_b[j]; + int nel_w = (int)(w->ne[0] * w->ne[1]); + std::vector w_data(nel_w), b_data(D); + sam3_read_f32(w, w_data.data(), nel_w); + sam3_read_f32(b, b_data.data(), D); + for (int o = 0; o < D; ++o) { + float sum = b_data[o]; + for (int i = 0; i < D; ++i) sum += w_data[o * D + i] * h[i]; + tmp[o] = (j < 2) ? std::max(0.0f, sum) : sum; + } + std::swap(h, tmp); + } - x = ggml_add(ctx, shortcut, ffn); + // Blend: obj_ptr = λ * projected + (1-λ) * no_obj_ptr + std::vector no_ptr(D); + ggml_backend_tensor_get(model.no_obj_ptr, no_ptr.data(), 0, D * sizeof(float)); + for (int i = 0; i < D; ++i) + out_ptr[i] = lambda * h[i] + (1.0f - lambda) * no_ptr[i]; + return; } - return x; -} + // SAM3 / default path: binary threshold + if (obj_score <= 0.0f) { + ggml_backend_tensor_get(model.no_obj_ptr, out_ptr, 0, D * sizeof(float)); + return; + } -// Build full fusion encoder graph (6 layers). -// image_feats: [D, N, B] where N=5184 (72*72), D=256 -// prompt_tokens: [D, T, B] text/exemplar features -// pos_enc: [D, N, B] sinusoidal positional encoding for image features -// Returns: conditioned_features [D, N, B] -static struct ggml_tensor* sam3_build_fenc_graph( - struct ggml_context* ctx, - const sam3_model& model, - struct ggml_tensor* image_feats, - struct ggml_tensor* prompt_tokens, - struct ggml_tensor* pos_enc, - struct ggml_tensor* prompt_attn_bias = nullptr) { - const auto& hp = model.hparams; - auto* x = image_feats; + std::vector h(D), tmp(D); + std::copy(sam_token_data, sam_token_data + D, h.data()); - for (int i = 0; i < hp.fenc_layers; ++i) { - x = sam3_fenc_layer_forward(ctx, model.fenc.layers[i], x, prompt_tokens, - pos_enc, prompt_attn_bias, hp.fenc_heads); - sam3_name_tensorf(x, "fenc_layer%d_out", i); - } + for (int j = 0; j < 3; ++j) { + auto* w = model.obj_ptr_proj_w[j]; + auto* b = model.obj_ptr_proj_b[j]; - return x; + int nel_w = (int)(w->ne[0] * w->ne[1]); + std::vector w_data(nel_w); + sam3_read_f32(w, w_data.data(), nel_w); + + std::vector b_data(D); + sam3_read_f32(b, b_data.data(), D); + + for (int o = 0; o < D; ++o) { + float sum = b_data[o]; + for (int i = 0; i < D; ++i) { + sum += w_data[o * D + i] * h[i]; + } + tmp[o] = (j < 2) ? std::max(0.0f, sum) : sum; + } + std::swap(h, tmp); + } + std::copy(h.begin(), h.end(), out_ptr); } /***************************************************************************** -** DETR decoder — graph building (6 layers) +** Tracker infrastructure (Phase 7, Step 7.4) *****************************************************************************/ -// inverse_sigmoid: log(x / (1 - x)), clamped to avoid inf -// Python reference uses eps=1e-3: x1 = x.clamp(min=eps), x2 = (1-x).clamp(min=eps) -static struct ggml_tensor* sam3_inverse_sigmoid(struct ggml_context* ctx, struct ggml_tensor* x) { - // clamp x to [1e-3, 1-1e-3] to match Python eps=1e-3 - x = ggml_clamp(ctx, x, 1e-3f, 1.0f - 1e-3f); - // log(x / (1 - x)) = log(x) - log(1 - x) - auto* log_x = ggml_log(ctx, x); - // Compute (1 - x) as (-1)*x + 1. We use ggml_scale_bias which takes float - // scalars (no tensor allocation needed, safe in no_alloc contexts). - auto* one_minus = ggml_scale_bias(ctx, x, -1.0f, 1.0f); - auto* log_1mx = ggml_log(ctx, one_minus); - return ggml_sub(ctx, log_x, log_1mx); +// Select memory frames for propagation (most recent + evenly spaced). +static std::vector sam3_select_memory_frames( + const std::vector& bank, + int max_slots) { + if ((int)bank.size() <= max_slots) { + std::vector all(bank.size()); + for (int i = 0; i < (int)bank.size(); ++i) all[i] = i; + return all; + } + std::vector selected; + selected.push_back(0); + selected.push_back((int)bank.size() - 1); + int remaining = max_slots - 2; + if (remaining > 0) { + float step = (float)(bank.size() - 2) / (remaining + 1); + for (int i = 0; i < remaining; ++i) { + int idx = 1 + (int)((i + 1) * step); + idx = std::min(idx, (int)bank.size() - 2); + selected.push_back(idx); + } + } + std::sort(selected.begin(), selected.end()); + selected.erase(std::unique(selected.begin(), selected.end()), selected.end()); + return selected; } -// Box refinement MLP (3 layers: D→D→D→4 with ReLU) -static struct ggml_tensor* sam3_bbox_mlp(struct ggml_context* ctx, - struct ggml_tensor* x, - struct ggml_tensor* w[3], - struct ggml_tensor* b[3]) { - for (int j = 0; j < 3; ++j) { - x = ggml_mul_mat(ctx, w[j], x); - x = ggml_add(ctx, x, b[j]); - if (j < 2) x = ggml_relu(ctx, x); +// Compute mask IoU between two binary masks. +static float sam3_mask_iou(const uint8_t* a, const uint8_t* b, int n) { + int inter = 0, uni = 0; + for (int i = 0; i < n; ++i) { + bool va = a[i] > 127; + bool vb = b[i] > 127; + if (va && vb) ++inter; + if (va || vb) ++uni; } - return x; + return (uni > 0) ? (float)inter / uni : 0.0f; } -// Build sinusoidal positional embedding for 4D reference points in the ggml graph. -// ref_boxes: [4, NQ, B] — (cx, cy, w, h) after sigmoid, B=1 -// sine_dim_t: [1, 64] — pre-computed angle multipliers (2π / 10000^(2i/128)) -// Returns: [512, NQ, B] sinusoidal embedding matching Python gen_sineembed_for_position -static struct ggml_tensor* sam3_build_sine_pos_embed_4d( - struct ggml_context* ctx, - struct ggml_tensor* ref_boxes, // [4, NQ, B] - struct ggml_tensor* sine_dim_t) { // [1, 64] - const int64_t NQ = ref_boxes->ne[1]; - - // Python output order: [cy, cx, w, h] → coord indices from boxes [cx(0),cy(1),w(2),h(3)] - const int coord_order[4] = {1, 0, 2, 3}; - - struct ggml_tensor* coord_embeds[4]; - - for (int c = 0; c < 4; ++c) { - int ci = coord_order[c]; - // Extract one coordinate: view into ref_boxes [4, NQ, 1] at element ci - auto* coord = ggml_view_2d(ctx, ref_boxes, 1, NQ, - ref_boxes->nb[1], ci * sizeof(float)); // [1, NQ] - - // Outer product: angles[i, q] = dim_t[i] * coord[q] - // ggml_mul_mat(A=[1,64], B=[1,NQ]) = A^T @ B = [64,1]@[1,NQ] = [64, NQ] - auto* angles = ggml_mul_mat(ctx, sine_dim_t, coord); // [64, NQ] +/***************************************************************************** +** Post-processing: hole filling and sprinkle removal (Phase 7, Step 7.8) +*****************************************************************************/ - auto* sin_vals = ggml_sin(ctx, angles); // [64, NQ] - auto* cos_vals = ggml_cos(ctx, angles); // [64, NQ] +// Fill small holes in a binary mask using BFS connected components. +static void sam3_fill_holes(uint8_t* mask, int w, int h, int area_threshold) { + const int n = w * h; + std::vector labels(n, -1); + int next_label = 0; + std::vector component_sizes; - // Interleave: [sin_0, cos_0, sin_1, cos_1, ...] - auto* sin_r = ggml_reshape_3d(ctx, sin_vals, 1, 64, NQ); - auto* cos_r = ggml_reshape_3d(ctx, cos_vals, 1, 64, NQ); - auto* interleaved = ggml_concat(ctx, sin_r, cos_r, 0); // [2, 64, NQ] - coord_embeds[c] = ggml_reshape_2d(ctx, interleaved, 128, NQ); + for (int i = 0; i < n; ++i) { + if (mask[i] > 127 || labels[i] >= 0) continue; + int label = next_label++; + component_sizes.push_back(0); + std::vector queue; + queue.push_back(i); + labels[i] = label; + int head = 0; + bool touches_border = false; + while (head < (int)queue.size()) { + int p = queue[head++]; + component_sizes[label]++; + int px = p % w, py = p / w; + if (px == 0 || px == w - 1 || py == 0 || py == h - 1) touches_border = true; + int dx[] = {-1, 1, 0, 0}; + int dy[] = {0, 0, -1, 1}; + for (int d = 0; d < 4; ++d) { + int nx2 = px + dx[d], ny2 = py + dy[d]; + if (nx2 < 0 || nx2 >= w || ny2 < 0 || ny2 >= h) continue; + int ni = ny2 * w + nx2; + if (mask[ni] <= 127 && labels[ni] < 0) { + labels[ni] = label; + queue.push_back(ni); + } + } + } + if (touches_border) component_sizes[label] = area_threshold + 1; + } + for (int i = 0; i < n; ++i) { + if (mask[i] <= 127 && labels[i] >= 0 && component_sizes[labels[i]] <= area_threshold) { + mask[i] = 255; + } } - - // Concatenate all 4 coordinates → [512, NQ] - auto* embed = ggml_concat(ctx, coord_embeds[0], coord_embeds[1], 0); // [256, NQ] - embed = ggml_concat(ctx, embed, coord_embeds[2], 0); // [384, NQ] - embed = ggml_concat(ctx, embed, coord_embeds[3], 0); // [512, NQ] - - return embed; } -// Build query positional encoding from reference boxes via sine embed + ref_point_head MLP. -// ref_boxes: [4, NQ, 1] — after sigmoid -// sine_dim_t: [1, 64] -// Returns: [D, NQ+1, 1] (zeros for presence token at index 0, MLP output for object queries) -static struct ggml_tensor* sam3_build_query_pos( - struct ggml_context* ctx, - const sam3_model& model, - struct ggml_tensor* ref_boxes, // [4, NQ, 1] - struct ggml_tensor* sine_dim_t, // [1, 64] - int layer_idx = -1) { - const auto& tensors = model.tensors; - const int64_t NQ = ref_boxes->ne[1]; - const int D = model.hparams.neck_dim; // 256 - - // 1. Sine positional embedding: [512, NQ] - auto* sine_embed = sam3_build_sine_pos_embed_4d(ctx, ref_boxes, sine_dim_t); - if (layer_idx == 0) { - ggml_set_name(sine_embed, "ddec_query_sine_0"); - } +// Remove small foreground sprinkles. +static void sam3_remove_sprinkles(uint8_t* mask, int w, int h, int area_threshold) { + const int n = w * h; + std::vector labels(n, -1); + int next_label = 0; + std::vector component_sizes; - // 2. ref_point_head MLP: 512 → 256 → 256 - // Layer 0: relu(W0 @ sine_embed + b0) - auto* h = ggml_mul_mat(ctx, tensors.at("ddec.ref_point_head.layers.0.weight"), sine_embed); - h = ggml_add(ctx, h, tensors.at("ddec.ref_point_head.layers.0.bias")); - h = ggml_relu(ctx, h); - // Layer 1: W1 @ h + b1 (no activation) - auto* qpos_obj = ggml_mul_mat(ctx, tensors.at("ddec.ref_point_head.layers.1.weight"), h); - qpos_obj = ggml_add(ctx, qpos_obj, tensors.at("ddec.ref_point_head.layers.1.bias")); - // qpos_obj: [D, NQ] - if (layer_idx == 0) { - ggml_set_name(qpos_obj, "ddec_query_pos_0"); + for (int i = 0; i < n; ++i) { + if (mask[i] <= 127 || labels[i] >= 0) continue; + int label = next_label++; + component_sizes.push_back(0); + std::vector queue; + queue.push_back(i); + labels[i] = label; + int head = 0; + while (head < (int)queue.size()) { + int p = queue[head++]; + component_sizes[label]++; + int px = p % w, py = p / w; + int dx[] = {-1, 1, 0, 0}; + int dy[] = {0, 0, -1, 1}; + for (int d = 0; d < 4; ++d) { + int nx2 = px + dx[d], ny2 = py + dy[d]; + if (nx2 < 0 || nx2 >= w || ny2 < 0 || ny2 >= h) continue; + int ni = ny2 * w + nx2; + if (mask[ni] > 127 && labels[ni] < 0) { + labels[ni] = label; + queue.push_back(ni); + } + } + } + } + for (int i = 0; i < n; ++i) { + if (mask[i] > 127 && labels[i] >= 0 && component_sizes[labels[i]] <= area_threshold) { + mask[i] = 0; + } } - - // 3. Reshape to 3D and prepend zeros for presence token - qpos_obj = ggml_reshape_3d(ctx, qpos_obj, D, NQ, 1); - auto* qpos_pres = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, 1, 1); - ggml_set_name(qpos_pres, "ddec_query_pos_pres"); - ggml_set_input(qpos_pres); // zeros — set by caller - - return ggml_concat(ctx, qpos_pres, qpos_obj, 1); // [D, NQ+1, 1] } -// Build box-relative positional bias for DETR cross-attention. -// ref_boxes: [4, N_q, B] — (cx, cy, w, h) in [0,1] -// rpb_coords: [feat_hw] — normalized coords [0/H, 1/H, ..., (H-1)/H] (input tensor) -// Returns: bias tensor [N_kv, N_q+1, n_heads, B] for ggml_flash_attn_ext mask -// -// Python _get_rpb_matrix with boxRPB="log": -// 1. boxes → xyxy -// 2. deltas_x[q,w,:2] = [coord_w - x0, coord_w - x1] -// 3. deltas_y[q,h,:2] = [coord_h - y0, coord_h - y1] -// 4. log transform: sign(d*8) * log2(|d*8|+1) / log2(8) -// 5. MLP: [2] → [256] → [n_heads] -// 6. outer sum: B[h,w] = delta_y[h] + delta_x[w] -static struct ggml_tensor* sam3_compute_box_rpb( - struct ggml_context* ctx, - const sam3_model& model, - struct ggml_tensor* ref_boxes, // [4, N_q, B] - struct ggml_tensor* rpb_coords, // [feat_hw] — pre-filled grid coordinates - int feat_hw, - int layer_idx = -1) { - const int64_t NQ = ref_boxes->ne[1]; - const int NH = model.hparams.ddec_heads; // 8 - const int W = feat_hw; - const int H = feat_hw; - const auto& tensors = model.tensors; +// Resolve overlapping masks: higher-scoring instances take priority. +static void sam3_resolve_overlaps(std::vector& dets) { + if (dets.size() <= 1) return; + const int w = dets[0].mask.width; + const int h = dets[0].mask.height; + if (w == 0 || h == 0) return; + std::sort(dets.begin(), dets.end(), [](const sam3_detection& a, const sam3_detection& b) { + return a.score > b.score; + }); + const int n = w * h; + for (int i = 0; i < n; ++i) { + bool claimed = false; + for (auto& det : dets) { + if (det.mask.data.empty()) continue; + if (claimed) { + det.mask.data[i] = 0; + } else if (det.mask.data[i] > 127) { + claimed = true; + } + } + } +} - // ── 1. Convert cxcywh → xyxy ───────────────────────────────────────── - // ggml_view_2d on strided data is non-contiguous — ggml_scale requires contiguous. - // Use ggml_cont to make each coordinate slice contiguous. - auto* cx = ggml_cont(ctx, ggml_view_2d(ctx, ref_boxes, 1, NQ, ref_boxes->nb[1], 0)); - auto* cy = ggml_cont(ctx, ggml_view_2d(ctx, ref_boxes, 1, NQ, ref_boxes->nb[1], 1 * sizeof(float))); - auto* bw = ggml_cont(ctx, ggml_view_2d(ctx, ref_boxes, 1, NQ, ref_boxes->nb[1], 2 * sizeof(float))); - auto* bh = ggml_cont(ctx, ggml_view_2d(ctx, ref_boxes, 1, NQ, ref_boxes->nb[1], 3 * sizeof(float))); - // x0 = cx - w/2, x1 = cx + w/2 - auto* half_w = ggml_scale(ctx, bw, 0.5f); - auto* half_h = ggml_scale(ctx, bh, 0.5f); - auto* x0 = ggml_sub(ctx, cx, half_w); // [1, NQ] - auto* x1 = ggml_add(ctx, cx, half_w); - auto* y0 = ggml_sub(ctx, cy, half_h); - auto* y1 = ggml_add(ctx, cy, half_h); +/***************************************************************************** +** Post-processing: NMS, bilinear interpolation, mask binarization +*****************************************************************************/ - // ── 2. Compute deltas via outer subtract ────────────────────────────── - // coords: [W] → reshape to [W, 1] for outer subtract - auto* cw = ggml_reshape_2d(ctx, rpb_coords, W, 1); // [W, 1] +// Compute IoU between two boxes [x0, y0, x1, y1]. +static float sam3_box_iou(const sam3_box& a, const sam3_box& b) { + float x0 = std::max(a.x0, b.x0); + float y0 = std::max(a.y0, b.y0); + float x1 = std::min(a.x1, b.x1); + float y1 = std::min(a.y1, b.y1); + float inter = std::max(0.0f, x1 - x0) * std::max(0.0f, y1 - y0); + float area_a = (a.x1 - a.x0) * (a.y1 - a.y0); + float area_b = (b.x1 - b.x0) * (b.y1 - b.y0); + float uni = area_a + area_b - inter; + return (uni > 0.0f) ? inter / uni : 0.0f; +} - // Outer subtract: delta[w, q] = coord[w] - edge[q] - // Use ggml_mul_mat trick: not applicable for subtraction. - // Instead: repeat coords to [W, NQ], repeat edge to [W, NQ], subtract. - auto* shape_wn = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, W, NQ); +// Non-maximum suppression on detections, sorted by score descending. +// Returns indices of kept detections. +static std::vector sam3_nms(const std::vector& dets, float iou_thresh) { + // Sort indices by score descending + std::vector indices(dets.size()); + for (int i = 0; i < (int)dets.size(); ++i) indices[i] = i; + std::sort(indices.begin(), indices.end(), [&](int a, int b) { + return dets[a].score > dets[b].score; + }); - auto* cw_rep = ggml_repeat(ctx, cw, shape_wn); // [W, NQ] (each column = coords) - auto* x0_t = ggml_cont(ctx, ggml_transpose(ctx, x0)); // [NQ, 1] - auto* x0_rep = ggml_repeat(ctx, ggml_reshape_2d(ctx, x0_t, 1, NQ), shape_wn); - auto* x1_t = ggml_cont(ctx, ggml_transpose(ctx, x1)); - auto* x1_rep = ggml_repeat(ctx, ggml_reshape_2d(ctx, x1_t, 1, NQ), shape_wn); - auto* y0_t = ggml_cont(ctx, ggml_transpose(ctx, y0)); - auto* y0_rep = ggml_repeat(ctx, ggml_reshape_2d(ctx, y0_t, 1, NQ), shape_wn); - auto* y1_t = ggml_cont(ctx, ggml_transpose(ctx, y1)); - auto* y1_rep = ggml_repeat(ctx, ggml_reshape_2d(ctx, y1_t, 1, NQ), shape_wn); + std::vector suppressed(dets.size(), false); + std::vector keep; - auto* dx0 = ggml_sub(ctx, cw_rep, x0_rep); // [W, NQ] - auto* dx1 = ggml_sub(ctx, cw_rep, x1_rep); - auto* dy0 = ggml_sub(ctx, cw_rep, y0_rep); // reusing coords for H (H==W) - auto* dy1 = ggml_sub(ctx, cw_rep, y1_rep); + for (int idx : indices) { + if (suppressed[idx]) continue; + keep.push_back(idx); + for (int j : indices) { + if (suppressed[j] || j == idx) continue; + if (sam3_box_iou(dets[idx].box, dets[j].box) > iou_thresh) { + suppressed[j] = true; + } + } + } - // Stack into [2, W, NQ]: reshape each to [1, W, NQ], concat dim 0 - auto* dx0_r = ggml_reshape_3d(ctx, dx0, 1, W, NQ); - auto* dx1_r = ggml_reshape_3d(ctx, dx1, 1, W, NQ); - auto* deltas_x = ggml_concat(ctx, dx0_r, dx1_r, 0); // [2, W, NQ] - auto* dy0_r = ggml_reshape_3d(ctx, dy0, 1, H, NQ); - auto* dy1_r = ggml_reshape_3d(ctx, dy1, 1, H, NQ); - auto* deltas_y = ggml_concat(ctx, dy0_r, dy1_r, 0); // [2, H, NQ] + return keep; +} - // ── 3. Log transform: sign(d*8) * log2(|d*8|+1) / log2(8) ──────────── - const float scale8 = 8.0f; - const float inv_log2_8 = 1.0f / log2f(8.0f); // = 1/3 +// Bilinear interpolation of a flat mask [H_in * W_in] to [H_out * W_out]. +// Uses double for coordinate math to match PyTorch F.interpolate precision. +static std::vector sam3_bilinear_interpolate(const float* src, int src_w, int src_h, + int dst_w, int dst_h) { + std::vector dst(dst_w * dst_h); + const double sx = (double)src_w / dst_w; + const double sy = (double)src_h / dst_h; - auto rpb_log = [&](struct ggml_tensor* d) -> struct ggml_tensor* { - auto* d8 = ggml_scale(ctx, d, scale8); - auto* sign_d = ggml_sgn(ctx, d8); - auto* abs_d = ggml_abs(ctx, d8); - auto* log_val = ggml_log(ctx, ggml_scale_bias(ctx, abs_d, 1.0f, 1.0f)); - // log2(x) = ln(x) / ln(2) - log_val = ggml_scale(ctx, log_val, 1.0f / logf(2.0f)); - return ggml_mul(ctx, sign_d, ggml_scale(ctx, log_val, inv_log2_8)); - }; + for (int y = 0; y < dst_h; ++y) { + double fy = (y + 0.5) * sy - 0.5; + fy = std::max(0.0, std::min(fy, (double)(src_h - 1))); + const int y0 = std::min((int)fy, src_h - 2); + const int y1 = y0 + 1; + const float wy = (float)(fy - y0); - deltas_x = rpb_log(deltas_x); // [2, W, NQ] - deltas_y = rpb_log(deltas_y); // [2, H, NQ] + for (int x = 0; x < dst_w; ++x) { + double fx = (x + 0.5) * sx - 0.5; + fx = std::max(0.0, std::min(fx, (double)(src_w - 1))); + const int x0 = std::min((int)fx, src_w - 2); + const int x1 = x0 + 1; + const float wx = (float)(fx - x0); - // ── 4. MLP: [2, W*NQ] → [NH, W*NQ] ─────────────────────────────────── - // boxRPB_embed_x: MLP(2, 256, 8, 2) = Linear(2→256)+ReLU, Linear(256→8) - // Reshape to [2, W*NQ] so matmul treats each (w, q) pair as a sample - auto rpb_mlp = [&](struct ggml_tensor* d, const char* axis) -> struct ggml_tensor* { - int64_t spatial = d->ne[1]; - int64_t nq = d->ne[2]; - auto* flat = ggml_reshape_2d(ctx, d, 2, spatial * nq); // [2, W*NQ] - auto wn0 = std::string("ddec.boxRPB_embed_") + axis + ".layers.0.weight"; - auto bn0 = std::string("ddec.boxRPB_embed_") + axis + ".layers.0.bias"; - auto wn1 = std::string("ddec.boxRPB_embed_") + axis + ".layers.1.weight"; - auto bn1 = std::string("ddec.boxRPB_embed_") + axis + ".layers.1.bias"; - flat = ggml_mul_mat(ctx, tensors.at(wn0), flat); - flat = ggml_add(ctx, flat, tensors.at(bn0)); - flat = ggml_relu(ctx, flat); - flat = ggml_mul_mat(ctx, tensors.at(wn1), flat); - flat = ggml_add(ctx, flat, tensors.at(bn1)); - // flat: [NH, W*NQ] → reshape to [NH, spatial, NQ] - return ggml_reshape_3d(ctx, flat, NH, spatial, nq); - }; + float v = (1 - wy) * ((1 - wx) * src[y0 * src_w + x0] + wx * src[y0 * src_w + x1]) + wy * ((1 - wx) * src[y1 * src_w + x0] + wx * src[y1 * src_w + x1]); + dst[y * dst_w + x] = v; + } + } + return dst; +} - auto* rpb_x = rpb_mlp(deltas_x, "x"); // [NH, W, NQ] - auto* rpb_y = rpb_mlp(deltas_y, "y"); // [NH, H, NQ] +// Convert (cx, cy, w, h) in [0,1] to (x0, y0, x1, y1) in pixel coordinates. +static sam3_box sam3_cxcywh_to_xyxy(float cx, float cy, float w, float h, + int img_w, int img_h) { + sam3_box box; + box.x0 = (cx - w * 0.5f) * img_w; + box.y0 = (cy - h * 0.5f) * img_h; + box.x1 = (cx + w * 0.5f) * img_w; + box.y1 = (cy + h * 0.5f) * img_h; + // Clamp to image bounds + box.x0 = std::max(0.0f, std::min(box.x0, (float)img_w)); + box.y0 = std::max(0.0f, std::min(box.y0, (float)img_h)); + box.x1 = std::max(0.0f, std::min(box.x1, (float)img_w)); + box.y1 = std::max(0.0f, std::min(box.y1, (float)img_h)); + return box; +} - // ── 5. Outer sum: B[nh, w, h, q] = rpb_y[nh, h, q] + rpb_x[nh, w, q] ─ - // Reshape for broadcasting: - // rpb_y → [NH, 1, H, NQ] - // rpb_x → [NH, W, 1, NQ] - // - // Keep W in ne[1] so reshaping [NH, W, H, NQ] → [NH, H*W, NQ, 1] - // preserves Python's flatten(H, W) order where W is the fast spatial axis. - auto* rpb_y_4d = ggml_reshape_4d(ctx, rpb_y, NH, 1, H, NQ); - auto* rpb_x_4d = ggml_reshape_4d(ctx, rpb_x, NH, W, 1, NQ); +/***************************************************************************** +** Image segmentation — PCS (text-prompted) +*****************************************************************************/ - // ggml_add broadcasts: where one dim is 1, the other is used - auto* rpb_hw = ggml_repeat(ctx, rpb_y_4d, - ggml_new_tensor_4d(ctx, GGML_TYPE_F32, NH, W, H, NQ)); - auto* rpb_hw_x = ggml_repeat(ctx, rpb_x_4d, - ggml_new_tensor_4d(ctx, GGML_TYPE_F32, NH, W, H, NQ)); - auto* rpb = ggml_add(ctx, rpb_hw, rpb_hw_x); // [NH, W, H, NQ] - - // ── 6. Reshape to [H*W, NQ, NH, 1] for flash_attn_ext mask ─────────── - // Current: [NH, W, H, NQ]. Need: [N_kv=H*W, NQ, NH, B=1] - rpb = ggml_reshape_4d(ctx, rpb, NH, H * W, NQ, 1); - rpb = ggml_cont(ctx, ggml_permute(ctx, rpb, 2, 0, 1, 3)); // [H*W, NQ, NH, 1] +sam3_result sam3_segment_pcs(sam3_state& state, + const sam3_model& model, + const sam3_pcs_params& params) { + if (model.hparams.visual_only || model.hparams.is_sam2()) { + fprintf(stderr, "%s: ERROR: PCS not available on %s model\n", + __func__, model.hparams.is_sam2() ? "SAM2" : "visual-only"); + return sam3_result{}; + } - // Prepend zeros for presence token: mask for presence token has no box-relative bias - auto* pres_mask = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, H * W, 1, NH, 1); - ggml_set_name(pres_mask, "rpb_pres_zeros"); - ggml_set_input(pres_mask); // zeros — set by caller +#if SAM3_LOG_LEVEL >= 1 + auto t_start = std::chrono::high_resolution_clock::now(); +#endif + const auto& hp = model.hparams; + const int D = hp.neck_dim; // 256 + const int H = hp.n_img_embd(); // 72 + const int L = hp.text_ctx_len; // 32 + const int NQ = hp.ddec_num_queries; // 200 + const int N_spatial = H * H; // 5184 + sam3_result result; - // [H*W, NQ+1, NH, 1] - auto* full_rpb = ggml_concat(ctx, pres_mask, rpb, 1); - full_rpb = ggml_cont(ctx, full_rpb); - if (layer_idx == 0) { - ggml_set_name(full_rpb, "ddec_rpb_mask_0"); + // ── Check that image has been encoded ──────────────────────────────── + if (!state.neck_det[0]) { + fprintf(stderr, "%s: image not encoded — call sam3_encode_image first\n", __func__); + return result; } - return full_rpb; -} - -// Single DETR decoder layer. -// queries: [D, N_q, B] where N_q = 201 (200 object queries + 1 presence token) -// query_pos: [D, N_q, B] positional encoding for queries -// enc_feats: [D, N_kv, B] conditioned image features from fusion encoder -// enc_pos: [D, N_kv, B] positional encoding for image features -// text_feats: [D, T, B] text features -// rpb_mask: [N_kv, N_q, n_heads, B] box-relative positional bias (or nullptr) -// Returns: updated queries [D, N_q, B] -static struct ggml_tensor* sam3_ddec_layer_forward( - struct ggml_context* ctx, - const sam3_ddec_layer& ly, - struct ggml_tensor* queries, - struct ggml_tensor* query_pos, - struct ggml_tensor* enc_feats, - struct ggml_tensor* enc_pos, - struct ggml_tensor* text_feats, - int n_heads, - struct ggml_tensor* text_attn_bias = nullptr, - struct ggml_tensor* rpb_mask = nullptr, - int layer_idx = -1) { - const int64_t D = queries->ne[0]; + // ── Tokenize text prompt ───────────────────────────────────────────── + auto token_ids = sam3_tokenize(const_cast(model.tokenizer), + params.text_prompt, L); + if (token_ids.empty()) { + fprintf(stderr, "%s: failed to tokenize text prompt\n", __func__); + return result; + } - // Python decoder layer order (all post-norm): - // 1. Self-attention → norm2 (post-norm) - // 2. Text cross-attention (ca_text) → catext_norm (post-norm) - // 3. Image cross-attention (cross_attn) → norm1 (post-norm) - // 4. FFN → norm3 (post-norm) - // - // Norm weight mapping: - // ly.norm2_w = ".norm2.weight" = Python norm2 (post-SA) - // ly.norm3_w = ".norm_ca_text.weight" = Python catext_norm (post-text-CA) - // ly.norm1_w = ".norm1.weight" = Python norm1 (post-image-CA) - // ly.norm4_w = ".norm3.weight" = Python norm3 (post-FFN) + SAM3_LOG(2, "%s: text='%s', %zu tokens\n", __func__, + params.text_prompt.c_str(), token_ids.size()); - // 1. Self-attention among queries (post-norm) - { - // Q = K = queries + query_pos, V = queries (no pos) - auto* q_in = ggml_add(ctx, queries, query_pos); + // ── Helper: run a sub-graph with its own context and allocator ────── + // Each stage below follows this exact pattern: + // 1. Create fresh ggml_context + graph + allocator + // 2. Create INPUT tensors (ggml_set_input) + // 3. Build the computation graph + // 4. Mark outputs (ggml_set_output) + // 5. Allocate → set input data → compute → read output data + // 6. Free allocator + context + // This ensures ZERO buffer sharing between stages. - auto* q_w = ggml_view_2d(ctx, ly.sa_in_proj_w, D, D, ly.sa_in_proj_w->nb[1], 0); - auto* k_w = ggml_view_2d(ctx, ly.sa_in_proj_w, D, D, ly.sa_in_proj_w->nb[1], D * ly.sa_in_proj_w->nb[1]); - auto* v_w = ggml_view_2d(ctx, ly.sa_in_proj_w, D, D, ly.sa_in_proj_w->nb[1], 2 * D * ly.sa_in_proj_w->nb[1]); - auto* q_b = ggml_view_1d(ctx, ly.sa_in_proj_b, D, 0); - auto* k_b = ggml_view_1d(ctx, ly.sa_in_proj_b, D, D * sizeof(float)); - auto* v_b = ggml_view_1d(ctx, ly.sa_in_proj_b, D, 2 * D * sizeof(float)); + // ── Pre-compute shared CPU data used by multiple stages ───────────── + const int n_boxes = (int)(params.pos_exemplars.size() + params.neg_exemplars.size()); + const int N_geo = n_boxes + 1; // +1 for CLS + const int T = L + N_geo; // total prompt tokens (text + geometry) - const int64_t N = queries->ne[1]; - const int64_t B = queries->ne[2]; - const int64_t HD = D / n_heads; + // Read image features and PE from state into CPU buffers (used by multiple stages) + std::vector img_feats_cpu(D * N_spatial); + std::vector img_pe_cpu(D * N_spatial); + ggml_backend_tensor_get(state.neck_det[2], img_feats_cpu.data(), 0, + D * N_spatial * sizeof(float)); + ggml_backend_tensor_get(state.neck_det_pe[2], img_pe_cpu.data(), 0, + D * N_spatial * sizeof(float)); - auto* Q = ggml_add(ctx, ggml_mul_mat(ctx, q_w, q_in), q_b); - auto* K = ggml_add(ctx, ggml_mul_mat(ctx, k_w, q_in), k_b); - auto* V = ggml_add(ctx, ggml_mul_mat(ctx, v_w, queries), v_b); + // Pre-compute constant input vectors + std::vector sine_dim_t_cpu(64); + for (int i = 0; i < 64; ++i) + sine_dim_t_cpu[i] = 2.0f * (float)M_PI / powf(10000.0f, 2.0f * (float)i / 128.0f); - Q = ggml_reshape_4d(ctx, Q, HD, n_heads, N, B); - Q = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 1, 3)); - K = ggml_reshape_4d(ctx, K, HD, n_heads, N, B); - K = ggml_cont(ctx, ggml_permute(ctx, K, 0, 2, 1, 3)); - V = ggml_reshape_4d(ctx, V, HD, n_heads, N, B); - V = ggml_permute(ctx, V, 0, 2, 1, 3); + std::vector rpb_coords_cpu(H); + for (int i = 0; i < H; ++i) rpb_coords_cpu[i] = (float)i / (float)H; - float scale = 1.0f / sqrtf((float)HD); - auto* sa_out = ggml_flash_attn_ext(ctx, Q, K, V, nullptr, scale, 0.0f, 0.0f); - sa_out = ggml_reshape_3d(ctx, sa_out, D, N, B); - sa_out = ggml_mul_mat(ctx, ly.sa_out_proj_w, sa_out); - sa_out = ggml_add(ctx, sa_out, ly.sa_out_proj_b); + // Build combined attention bias for the prompt: [T] = [L + N_geo] + // 0.0 for valid tokens, -1e9 for padding text tokens, 0.0 for geo tokens + std::vector combined_bias_cpu(T); + for (int i = 0; i < L; ++i) + combined_bias_cpu[i] = (token_ids[i] != 0) ? 0.0f : -1.0e9f; + for (int i = 0; i < N_geo; ++i) + combined_bias_cpu[L + i] = 0.0f; - queries = ggml_add(ctx, queries, sa_out); - // Post-norm: norm2 (Python's post-SA norm) - queries = sam3_layer_norm(ctx, queries, ly.norm2_w, ly.norm2_b); - if (layer_idx == 0) { - ggml_set_name(queries, "ddec_layer0_after_sa"); - } - } + // Build text validity mask for DotProductScoring: [T] + int n_valid_tokens = 0; + for (int i = 0; i < L; ++i) + if (token_ids[i] != 0) ++n_valid_tokens; + n_valid_tokens += N_geo; + if (n_valid_tokens == 0) n_valid_tokens = 1; + float tvm_scale = (float)T / (float)n_valid_tokens; + std::vector text_valid_mask_cpu(T); + for (int i = 0; i < L; ++i) + text_valid_mask_cpu[i] = (token_ids[i] != 0) ? tvm_scale : 0.0f; + for (int i = 0; i < N_geo; ++i) + text_valid_mask_cpu[L + i] = tvm_scale; - // 2. Cross-attention to text tokens (post-norm) + /* + ** ── SUB-GRAPH 1: Text Encoder ──────────────────────────────────── + */ + std::vector text_feats_cpu(D * L); { - // Q = queries + query_pos, K = V = text_feats - auto* q_in = ggml_add(ctx, queries, query_pos); + const size_t sz = ggml_tensor_overhead() * 16384 + ggml_graph_overhead() * 2; + struct ggml_init_params gp = {sz, nullptr, true}; + auto* ctx = ggml_init(gp); - auto* q_w = ggml_view_2d(ctx, ly.ca_text_q_w, D, D, ly.ca_text_q_w->nb[1], 0); - auto* k_w = ggml_view_2d(ctx, ly.ca_text_q_w, D, D, ly.ca_text_q_w->nb[1], D * ly.ca_text_q_w->nb[1]); - auto* v_w = ggml_view_2d(ctx, ly.ca_text_q_w, D, D, ly.ca_text_q_w->nb[1], 2 * D * ly.ca_text_q_w->nb[1]); - auto* q_b = ggml_view_1d(ctx, ly.ca_text_q_b, D, 0); - auto* k_b = ggml_view_1d(ctx, ly.ca_text_q_b, D, D * sizeof(float)); - auto* v_b = ggml_view_1d(ctx, ly.ca_text_q_b, D, 2 * D * sizeof(float)); + auto* inp = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, L); + ggml_set_name(inp, "text_token_ids"); + ggml_set_input(inp); - const int64_t N_q = queries->ne[1]; - const int64_t N_kv = text_feats->ne[1]; - const int64_t B = queries->ne[2]; - const int64_t HD = D / n_heads; + auto* out = sam3_build_text_encoder_graph(ctx, inp, model); + ggml_set_output(out); - auto* Q = ggml_add(ctx, ggml_mul_mat(ctx, q_w, q_in), q_b); - auto* K = ggml_add(ctx, ggml_mul_mat(ctx, k_w, text_feats), k_b); - auto* V = ggml_add(ctx, ggml_mul_mat(ctx, v_w, text_feats), v_b); + auto* graph = ggml_new_graph_custom(ctx, 16384, false); + ggml_build_forward_expand(graph, out); - Q = ggml_reshape_4d(ctx, Q, HD, n_heads, N_q, B); - Q = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 1, 3)); - K = ggml_reshape_4d(ctx, K, HD, n_heads, N_kv, B); - K = ggml_cont(ctx, ggml_permute(ctx, K, 0, 2, 1, 3)); - V = ggml_reshape_4d(ctx, V, HD, n_heads, N_kv, B); - V = ggml_permute(ctx, V, 0, 2, 1, 3); + auto* causal = ggml_get_tensor(ctx, "causal_mask"); + auto* alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); + if (!ggml_gallocr_reserve(alloc, graph) || !ggml_gallocr_alloc_graph(alloc, graph)) { + fprintf(stderr, "%s: text encoder alloc failed\n", __func__); + ggml_gallocr_free(alloc); + ggml_free(ctx); + return result; + } - auto* text_mask = sam3_expand_token_attn_bias(ctx, text_attn_bias, N_q, n_heads, B); - float scale = 1.0f / sqrtf((float)HD); - auto* ca_out = ggml_flash_attn_ext(ctx, Q, K, V, text_mask, scale, 0.0f, 0.0f); - ca_out = ggml_reshape_3d(ctx, ca_out, D, N_q, B); - ca_out = ggml_mul_mat(ctx, ly.ca_text_out_w, ca_out); - ca_out = ggml_add(ctx, ca_out, ly.ca_text_out_b); + if (inp->buffer) { + ggml_backend_tensor_set(inp, token_ids.data(), 0, L * sizeof(int32_t)); + } else { + fprintf(stderr, "%s: ERROR: text encoder input tensor has no buffer!\n", __func__); + ggml_gallocr_free(alloc); + ggml_free(ctx); + return result; + } + if (causal && causal->buffer) { + std::vector cm(L * L); + sam3_fill_causal_mask(cm.data(), L); + ggml_backend_tensor_set(causal, cm.data(), 0, L * L * sizeof(ggml_fp16_t)); + } - queries = ggml_add(ctx, queries, ca_out); - // Post-norm: catext_norm (Python's post-text-CA norm) - queries = sam3_layer_norm(ctx, queries, ly.norm3_w, ly.norm3_b); - if (layer_idx == 0) { - ggml_set_name(queries, "ddec_layer0_after_text_ca"); + if (!sam3_graph_compute(model.backend, graph, state.n_threads)) { + ggml_gallocr_free(alloc); + ggml_free(ctx); + return result; } + ggml_backend_tensor_get(out, text_feats_cpu.data(), 0, D * L * sizeof(float)); + + ggml_gallocr_free(alloc); + ggml_free(ctx); } - // 3. Cross-attention to conditioned image features (post-norm) + SAM3_LOG(2, "%s: text encoder done\n", __func__); + + /* + ** ── SUB-GRAPH 2: Geometry Encoder ──────────────────────────────── + */ + std::vector geo_feats_cpu(D * N_geo); { - // Q = queries + query_pos, K = enc_feats + enc_pos, V = enc_feats - auto* q_in = ggml_add(ctx, queries, query_pos); - auto* k_in = ggml_add(ctx, enc_feats, enc_pos); + const size_t sz = ggml_tensor_overhead() * 4096 + ggml_graph_overhead() * 2; + struct ggml_init_params gp = {sz, nullptr, true}; + auto* ctx = ggml_init(gp); - auto* q_w = ggml_view_2d(ctx, ly.ca_q_w, D, D, ly.ca_q_w->nb[1], 0); - auto* k_w = ggml_view_2d(ctx, ly.ca_q_w, D, D, ly.ca_q_w->nb[1], D * ly.ca_q_w->nb[1]); - auto* v_w = ggml_view_2d(ctx, ly.ca_q_w, D, D, ly.ca_q_w->nb[1], 2 * D * ly.ca_q_w->nb[1]); - auto* q_b = ggml_view_1d(ctx, ly.ca_q_b, D, 0); - auto* k_b = ggml_view_1d(ctx, ly.ca_q_b, D, D * sizeof(float)); - auto* v_b = ggml_view_1d(ctx, ly.ca_q_b, D, 2 * D * sizeof(float)); - - const int64_t N_q = queries->ne[1]; - const int64_t N_kv = enc_feats->ne[1]; - const int64_t B = queries->ne[2]; - const int64_t HD = D / n_heads; - - auto* Q = ggml_add(ctx, ggml_mul_mat(ctx, q_w, q_in), q_b); - auto* K = ggml_add(ctx, ggml_mul_mat(ctx, k_w, k_in), k_b); - auto* V = ggml_add(ctx, ggml_mul_mat(ctx, v_w, enc_feats), v_b); - - Q = ggml_reshape_4d(ctx, Q, HD, n_heads, N_q, B); - Q = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 1, 3)); - K = ggml_reshape_4d(ctx, K, HD, n_heads, N_kv, B); - K = ggml_cont(ctx, ggml_permute(ctx, K, 0, 2, 1, 3)); - V = ggml_reshape_4d(ctx, V, HD, n_heads, N_kv, B); - V = ggml_permute(ctx, V, 0, 2, 1, 3); + auto* g_img = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, N_spatial, 1); + ggml_set_name(g_img, "geo_img"); + ggml_set_input(g_img); + auto* g_pe = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, N_spatial, 1); + ggml_set_name(g_pe, "geo_pe"); + ggml_set_input(g_pe); - float scale = 1.0f / sqrtf((float)HD); - struct ggml_tensor* ca_out = nullptr; + auto gr = sam3_build_geom_enc_graph(ctx, model, params, g_img, g_pe); + ggml_set_output(gr.geo_feats); - if (rpb_mask) { - // Keep the box-relative positional bias in fp32. The CPU flash-attn - // kernel reads mask storage as fp16, which is good enough for token - // padding masks but introduces avoidable drift here. - auto* kq = ggml_mul_mat(ctx, K, Q); // [N_kv, N_q, NH, B] - kq = ggml_soft_max_ext(ctx, kq, rpb_mask, scale, 0.0f); // [N_kv, N_q, NH, B] + auto* graph = ggml_new_graph_custom(ctx, 4096, false); + ggml_build_forward_expand(graph, gr.geo_feats); - auto* v_t = ggml_cont(ctx, ggml_transpose(ctx, V)); // [N_kv, HD, NH, B] - ca_out = ggml_mul_mat(ctx, v_t, kq); // [HD, N_q, NH, B] - ca_out = ggml_cont(ctx, ggml_permute(ctx, ca_out, 0, 2, 1, 3)); - } else { - ca_out = ggml_flash_attn_ext(ctx, Q, K, V, nullptr, scale, 0.0f, 0.0f); + auto* alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); + if (!ggml_gallocr_reserve(alloc, graph) || !ggml_gallocr_alloc_graph(alloc, graph)) { + fprintf(stderr, "%s: geometry encoder alloc failed\n", __func__); + ggml_gallocr_free(alloc); + ggml_free(ctx); + return result; } - ca_out = ggml_reshape_3d(ctx, ca_out, D, N_q, B); - ca_out = ggml_mul_mat(ctx, ly.ca_out_w, ca_out); - ca_out = ggml_add(ctx, ca_out, ly.ca_out_b); + ggml_backend_tensor_set(g_img, img_feats_cpu.data(), 0, D * N_spatial * sizeof(float)); + ggml_backend_tensor_set(g_pe, img_pe_cpu.data(), 0, D * N_spatial * sizeof(float)); - queries = ggml_add(ctx, queries, ca_out); - // Post-norm: norm1 (Python's post-image-CA norm) - queries = sam3_layer_norm(ctx, queries, ly.norm1_w, ly.norm1_b); - if (layer_idx == 0) { - ggml_set_name(queries, "ddec_layer0_after_img_ca"); + // Pre-transformer input (CLS + optional box embeddings) + { + const float* feats_ptr = n_boxes > 0 ? img_feats_cpu.data() : nullptr; + auto geom_data = sam3_precompute_geom_input(model, params, feats_ptr, H, H); + auto* gi = ggml_get_tensor(ctx, "geom_post_final_proj"); + if (gi) ggml_backend_tensor_set(gi, geom_data.data(), 0, geom_data.size() * sizeof(float)); } - } - - // 4. FFN (post-norm, ReLU) - { - auto* ffn = ggml_mul_mat(ctx, ly.ffn_fc1_w, queries); - ffn = ggml_add(ctx, ffn, ly.ffn_fc1_b); - ffn = ggml_relu(ctx, ffn); - ffn = ggml_mul_mat(ctx, ly.ffn_fc2_w, ffn); - ffn = ggml_add(ctx, ffn, ly.ffn_fc2_b); - queries = ggml_add(ctx, queries, ffn); - // Post-norm: norm3 (Python's post-FFN norm) - queries = sam3_layer_norm(ctx, queries, ly.norm4_w, ly.norm4_b); - if (layer_idx == 0) { - ggml_set_name(queries, "ddec_layer0_full_out"); + if (!sam3_graph_compute(model.backend, graph, state.n_threads)) { + ggml_gallocr_free(alloc); + ggml_free(ctx); + return result; } - } + ggml_backend_tensor_get(gr.geo_feats, geo_feats_cpu.data(), 0, + D * N_geo * sizeof(float)); - return queries; -} + ggml_gallocr_free(alloc); + ggml_free(ctx); + } -// DotProductScoring: classify queries against text features via dot product. -// -// Python reference (DotProductScoring.forward): -// 1. prompt_mlp(prompt) → residual MLP + LN on text features -// 2. mean_pool_text(result, prompt_mask) → pooled [BS, D] (only valid tokens) -// 3. prompt_proj(pooled) → [BS, D] -// 4. hs_proj(hs) → [num_layer, BS, N_q, D] -// 5. matmul(proj_hs, proj_pooled.unsqueeze(-1)) → dot product → [num_layer, BS, N_q, 1] -// 6. scale by 1/sqrt(D) -// 7. clamp to [-12, 12] -// -// query_outputs: [D, N_q, B] — the 200 object query outputs -// text_features: [D, T, B] — text encoder output (already through resizer) -// text_valid_mask: [T, 1, B] — 1.0 for valid tokens, 0.0 for padding (or nullptr for all-valid) -// Returns: class_scores [N_q, B] (one score per query per batch) -static struct ggml_tensor* sam3_dot_product_scoring( - struct ggml_context* ctx, - const sam3_model& model, - struct ggml_tensor* query_outputs, // [D, N_q, B] - struct ggml_tensor* text_features, // [D, T, B] - struct ggml_tensor* text_valid_mask) // [T, 1, B] or nullptr -{ - const auto& tensors = model.tensors; - const int64_t D = query_outputs->ne[0]; // 256 - const int64_t T = text_features->ne[1]; - const int64_t B = text_features->ne[2]; + SAM3_LOG(2, "%s: geometry encoder done\n", __func__); - // Step 1: Apply prompt_mlp on text features (residual MLP + LayerNorm) - auto* text_mlp = text_features; // [D, T, B] - auto* orig_text = text_features; - text_mlp = ggml_mul_mat(ctx, tensors.at("scoring.prompt_mlp.layers.0.weight"), text_mlp); - text_mlp = ggml_add(ctx, text_mlp, tensors.at("scoring.prompt_mlp.layers.0.bias")); - text_mlp = ggml_relu(ctx, text_mlp); - text_mlp = ggml_mul_mat(ctx, tensors.at("scoring.prompt_mlp.layers.1.weight"), text_mlp); - text_mlp = ggml_add(ctx, text_mlp, tensors.at("scoring.prompt_mlp.layers.1.bias")); - text_mlp = ggml_add(ctx, text_mlp, orig_text); - text_mlp = sam3_layer_norm(ctx, text_mlp, - tensors.at("scoring.prompt_mlp.out_norm.weight"), - tensors.at("scoring.prompt_mlp.out_norm.bias")); - // text_mlp: [D, T, B] - ggml_set_name(text_mlp, "scoring_prompt_mlp_out"); + // Build combined prompt on CPU: [text_feats, geo_feats] → [D * T] + std::vector combined_prompt_cpu(D * T); + memcpy(combined_prompt_cpu.data(), text_feats_cpu.data(), D * L * sizeof(float)); + memcpy(combined_prompt_cpu.data() + D * L, geo_feats_cpu.data(), D * N_geo * sizeof(float)); - // Step 2: Mean-pool over valid text tokens → [D, 1, B] - // Python: pooled = (prompt * is_valid).sum(0) / num_valid - struct ggml_tensor* text_pooled; - if (text_valid_mask) { - // Permute to [T, D, B] for masked pooling - auto* tp = ggml_cont(ctx, ggml_permute(ctx, text_mlp, 1, 0, 2, 3)); // [T, D, B] - // Mask: text_valid_mask is [T, 1, B] — broadcast multiply zeros out padding - tp = ggml_mul(ctx, tp, text_valid_mask); // [T, D, B] with padding zeroed - // Sum over T dimension: pool_1d with SUM kernel=T - // ggml_pool_1d AVG divides by T; we want SUM then divide by n_valid. - // Use AVG and then scale by T/n_valid? Or use a manual approach. - // Simpler: sum via pool_1d with AVG, then scale by T/n_valid. - // But n_valid is dynamic. Instead: sum = mean * T, then divide by n_valid. - // We pass n_valid as part of the mask: text_valid_mask sums to n_valid. - // pool_1d(masked, AVG, T, T, 0) = sum(masked) / T. Multiply by T → sum(masked). - // Then divide by n_valid. But n_valid is a scalar we know CPU-side. - // For simplicity: compute AVG over ALL T positions (with padding zeroed out). - // This gives sum(valid) / T. To get sum(valid) / n_valid, scale by T / n_valid. - // We embed the scale factor into the mask: mask = (T / n_valid) for valid, 0 for pad. - // Then AVG(mask * features) = sum(valid * T/n_valid) / T = sum(valid) / n_valid. ✓ - // Caller should set mask values to T/n_valid for valid tokens, 0 for padding. - auto* pooled_t = ggml_pool_1d(ctx, tp, GGML_OP_POOL_AVG, (int)T, (int)T, 0); - text_pooled = ggml_cont(ctx, ggml_permute(ctx, pooled_t, 1, 0, 2, 3)); // [D, 1, B] - } else { - // All tokens valid — simple mean - auto* tp = ggml_cont(ctx, ggml_permute(ctx, text_mlp, 1, 0, 2, 3)); - auto* pooled_t = ggml_pool_1d(ctx, tp, GGML_OP_POOL_AVG, (int)T, (int)T, 0); - text_pooled = ggml_cont(ctx, ggml_permute(ctx, pooled_t, 1, 0, 2, 3)); - } - ggml_set_name(text_pooled, "scoring_pooled"); + SAM3_LOG(2, "%s: starting fusion encoder\n", __func__); + /* + ** ── SUB-GRAPH 3: Fusion Encoder ────────────────────────────────── + */ + std::vector fenc_output_cpu(D * N_spatial); + { + const size_t sz = ggml_tensor_overhead() * 16384 + ggml_graph_overhead() * 2; + struct ggml_init_params gp = {sz, nullptr, true}; + auto* ctx = ggml_init(gp); - // Step 3: Project pooled prompt through prompt_proj: D→D - auto* proj_pooled = ggml_mul_mat(ctx, tensors.at("scoring.prompt_proj.weight"), text_pooled); - proj_pooled = ggml_add(ctx, proj_pooled, tensors.at("scoring.prompt_proj.bias")); - // proj_pooled: [D, 1, B] - ggml_set_name(proj_pooled, "scoring_proj_pooled"); + auto* img = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, N_spatial, 1); + ggml_set_name(img, "fenc_img"); + ggml_set_input(img); + auto* pe = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, N_spatial, 1); + ggml_set_name(pe, "fenc_pe"); + ggml_set_input(pe); + auto* prompt = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, T, 1); + ggml_set_name(prompt, "fenc_prompt"); + ggml_set_input(prompt); + auto* bias = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, T, 1, 1); + ggml_set_name(bias, "fenc_bias"); + ggml_set_input(bias); - // Step 4: Project queries through hs_proj: D→D - auto* proj_hs = ggml_mul_mat(ctx, tensors.at("scoring.hs_proj.weight"), query_outputs); - proj_hs = ggml_add(ctx, proj_hs, tensors.at("scoring.hs_proj.bias")); - // proj_hs: [D, N_q, B] - ggml_set_name(proj_hs, "scoring_proj_hs"); + auto* out = sam3_build_fenc_graph(ctx, model, img, prompt, pe, bias); + ggml_set_output(out); - // Step 5: Dot product — for each query, dot with pooled prompt - // matmul(proj_hs, proj_pooled.unsqueeze(-1)) in Python = batched vector-matrix multiply - // ggml_mul_mat(A, B) = A^T @ B - // With A = proj_pooled [D, 1, B], B = proj_hs [D, N_q, B]: - // result = [1, N_q, B] — each element is dot product of query with pooled prompt - auto* scores = ggml_mul_mat(ctx, proj_pooled, proj_hs); // [1, N_q, B] + auto* graph = ggml_new_graph_custom(ctx, 16384, false); + ggml_build_forward_expand(graph, out); - // Step 6: Scale by 1/sqrt(D) - float scale = 1.0f / sqrtf((float)D); - scores = ggml_scale(ctx, scores, scale); + auto* alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); + if (!ggml_gallocr_reserve(alloc, graph) || !ggml_gallocr_alloc_graph(alloc, graph)) { + fprintf(stderr, "%s: fusion encoder alloc failed\n", __func__); + ggml_gallocr_free(alloc); + ggml_free(ctx); + return result; + } - // Step 7: Clamp to [-12, 12] - scores = ggml_clamp(ctx, scores, -12.0f, 12.0f); + ggml_backend_tensor_set(img, img_feats_cpu.data(), 0, D * N_spatial * sizeof(float)); + ggml_backend_tensor_set(pe, img_pe_cpu.data(), 0, D * N_spatial * sizeof(float)); + ggml_backend_tensor_set(prompt, combined_prompt_cpu.data(), 0, D * T * sizeof(float)); + ggml_backend_tensor_set(bias, combined_bias_cpu.data(), 0, T * sizeof(float)); - // Reshape to [N_q, B] - const int64_t N_q = query_outputs->ne[1]; - scores = ggml_reshape_2d(ctx, scores, N_q, B); - ggml_set_name(scores, "scoring_class_scores"); + if (!sam3_graph_compute(model.backend, graph, state.n_threads)) { + ggml_gallocr_free(alloc); + ggml_free(ctx); + return result; + } + ggml_backend_tensor_get(out, fenc_output_cpu.data(), 0, D * N_spatial * sizeof(float)); - return scores; -} + SAM3_LOG(2, "%s: fenc_out[0..4] = [%.6f, %.6f, %.6f, %.6f, %.6f]\n", + __func__, fenc_output_cpu[0], fenc_output_cpu[1], fenc_output_cpu[2], + fenc_output_cpu[3], fenc_output_cpu[4]); -// Build full DETR decoder graph. -// enc_feats: [D, N_kv, B] conditioned features from fusion encoder (N_kv=5184) -// enc_pos: [D, N_kv, B] positional encoding -// text_feats: [D, T, B] text features -// Returns struct with: -// queries: [D, 201, B] (all query outputs including presence token) -// pred_boxes: [4, 200, B] (cx, cy, w, h in [0,1]) -// class_scores: [200, B] -// presence_score: [1, B] -struct sam3_ddec_output { - struct ggml_tensor* queries; // [D, 201, B] - struct ggml_tensor* presence_feats; // [D, 1, B] pre-decoder-norm presence token - struct ggml_tensor* pred_boxes; // [4, 200, B] - struct ggml_tensor* class_scores; // [200, B] - struct ggml_tensor* presence_score; // [1, B] -}; + ggml_gallocr_free(alloc); + ggml_free(ctx); + } -static sam3_ddec_output sam3_build_ddec_graph( - struct ggml_context* ctx, - const sam3_model& model, - struct ggml_tensor* enc_feats, // [D, N_kv, B] - struct ggml_tensor* enc_pos, // [D, N_kv, B] - struct ggml_tensor* text_feats, // [D, T, B] - struct ggml_tensor* sine_dim_t, // [1, 64] — pre-computed angle multipliers - struct ggml_tensor* rpb_coords, // [feat_hw] — normalized grid coords (or nullptr) - struct ggml_tensor* text_attn_bias = nullptr, // [T, 1, B] additive text padding bias - struct ggml_tensor* text_valid_mask = nullptr) // [T, 1, B] for scoring (or nullptr) -{ - const auto& hp = model.hparams; - const auto& tensors = model.tensors; - const int D = hp.neck_dim; // 256 - const int NQ = hp.ddec_num_queries; // 200 - const int B = (int)enc_feats->ne[2]; // batch (1) - const int feat_hw = hp.n_img_embd(); // 72 + SAM3_LOG(2, "%s: fusion encoder done\n", __func__); + /* + ** ── SUB-GRAPH 4: DETR Decoder + Scoring ────────────────────────── + */ + std::vector scores_data(NQ); + std::vector boxes_data(4 * NQ); + std::vector queries_data(D * (NQ + 1)); // 201 = NQ + presence token + float presence_logit = 0.0f; + { + const size_t sz = ggml_tensor_overhead() * 32768 + ggml_graph_overhead() * 2; + struct ggml_init_params gp = {sz, nullptr, true}; + auto* ctx = ggml_init(gp); - // ── Initialize queries from query_embed ────────────────────────────── - auto* content = ggml_reshape_3d(ctx, model.ddec.query_embed, D, NQ, 1); - auto* pres_tok = ggml_reshape_3d(ctx, model.ddec.presence_token, D, 1, 1); - auto* queries = ggml_concat(ctx, pres_tok, content, 1); // [D, NQ+1, B=1] + auto* enc = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, N_spatial, 1); + ggml_set_name(enc, "ddec_enc"); + ggml_set_input(enc); + auto* pe = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, N_spatial, 1); + ggml_set_name(pe, "ddec_pe"); + ggml_set_input(pe); + auto* txt = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, T, 1); + ggml_set_name(txt, "ddec_text"); + ggml_set_input(txt); + auto* sdt = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, 64); + ggml_set_name(sdt, "sine_dim_t"); + ggml_set_input(sdt); + auto* rpb = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, H); + ggml_set_name(rpb, "rpb_coords"); + ggml_set_input(rpb); + auto* tab = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, T, 1, 1); + ggml_set_name(tab, "text_attn_bias"); + ggml_set_input(tab); + auto* tvm = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, T, 1, 1); + ggml_set_name(tvm, "text_valid_mask"); + ggml_set_input(tvm); - // Reference points: sigmoid → initial anchor boxes - auto* ref_pts_raw = ggml_cont(ctx, tensors.at("ddec.reference_points.weight")); // [4, NQ] - auto* ref_boxes = ggml_sigmoid(ctx, ref_pts_raw); // [4, NQ] - ref_boxes = ggml_reshape_3d(ctx, ref_boxes, 4, NQ, 1); // [4, NQ, 1] -#ifndef NDEBUG - auto* ref_boxes_dbg = ggml_cont(ctx, ref_boxes); - ggml_set_name(ref_boxes_dbg, "ddec_ref_boxes_init"); -#endif + auto dout = sam3_build_ddec_graph(ctx, model, enc, pe, txt, sdt, rpb, tab, tvm); + ggml_set_output(dout.class_scores); + ggml_set_output(dout.pred_boxes); + ggml_set_output(dout.presence_score); + ggml_set_output(dout.queries); - // ── Run decoder layers ─────────────────────────────────────────────── - // Per-layer: recompute query_pos from updated ref_boxes (matching Python exactly) - struct ggml_tensor* last_presence = pres_tok; - for (int i = 0; i < hp.ddec_layers; ++i) { - // Recompute query_pos from current ref_boxes via sine embed + ref_point_head MLP - auto* query_pos = sam3_build_query_pos(ctx, model, ref_boxes, sine_dim_t, i); + auto* graph = ggml_new_graph_custom(ctx, 65536, false); + ggml_build_forward_expand(graph, dout.class_scores); + ggml_build_forward_expand(graph, dout.pred_boxes); + ggml_build_forward_expand(graph, dout.presence_score); + ggml_build_forward_expand(graph, dout.queries); - // Compute box-relative positional bias for image cross-attention - struct ggml_tensor* rpb_mask = nullptr; - if (rpb_coords) { - rpb_mask = sam3_compute_box_rpb(ctx, model, ref_boxes, rpb_coords, feat_hw, i); + auto* alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); + if (!ggml_gallocr_reserve(alloc, graph) || !ggml_gallocr_alloc_graph(alloc, graph)) { + fprintf(stderr, "%s: DETR decoder alloc failed\n", __func__); + ggml_gallocr_free(alloc); + ggml_free(ctx); + return result; } - queries = sam3_ddec_layer_forward(ctx, model.ddec.layers[i], - queries, query_pos, - enc_feats, enc_pos, - text_feats, hp.ddec_heads, - text_attn_bias, - rpb_mask, - i); - - // Box refinement after each layer (on object queries only, not presence token) - auto* obj_q = ggml_view_3d(ctx, queries, D, NQ, 1, - queries->nb[1], queries->nb[2], 1 * queries->nb[1]); - obj_q = ggml_cont(ctx, obj_q); - sam3_name_tensorf(obj_q, "ddec_layer%d_out", i); + ggml_backend_tensor_set(enc, fenc_output_cpu.data(), 0, D * N_spatial * sizeof(float)); + ggml_backend_tensor_set(pe, img_pe_cpu.data(), 0, D * N_spatial * sizeof(float)); + ggml_backend_tensor_set(txt, combined_prompt_cpu.data(), 0, D * T * sizeof(float)); + ggml_backend_tensor_set(sdt, sine_dim_t_cpu.data(), 0, 64 * sizeof(float)); + ggml_backend_tensor_set(rpb, rpb_coords_cpu.data(), 0, H * sizeof(float)); + ggml_backend_tensor_set(tab, combined_bias_cpu.data(), 0, T * sizeof(float)); + ggml_backend_tensor_set(tvm, text_valid_mask_cpu.data(), 0, T * sizeof(float)); - auto* pres_q = ggml_view_3d(ctx, queries, D, 1, 1, - queries->nb[1], queries->nb[2], 0); - pres_q = ggml_cont(ctx, pres_q); - if (i == 0) { - ggml_set_name(pres_q, "ddec_layer0_presence"); + // Presence token position encoding: zeros + auto* qpos = ggml_get_tensor(ctx, "ddec_query_pos_pres"); + if (qpos) { + std::vector z(D, 0.0f); + ggml_backend_tensor_set(qpos, z.data(), 0, D * sizeof(float)); + } + auto* rpbz = ggml_get_tensor(ctx, "rpb_pres_zeros"); + if (rpbz) { + int n = (int)(rpbz->ne[0] * rpbz->ne[1] * rpbz->ne[2] * rpbz->ne[3]); + std::vector z(n, 0.0f); + ggml_backend_tensor_set(rpbz, z.data(), 0, n * sizeof(float)); } - last_presence = pres_q; - - // Apply the final decoder norm before box refinement (use_normed_output_consistently) - auto* obj_q_normed = sam3_layer_norm(ctx, obj_q, - tensors.at("ddec.norm.weight"), - tensors.at("ddec.norm.bias")); - // Shared bbox_embed MLP - auto* bd = obj_q_normed; - for (int j = 0; j < 3; ++j) { - auto wn = "ddec.bbox_embed.layers." + std::to_string(j) + ".weight"; - auto bn = "ddec.bbox_embed.layers." + std::to_string(j) + ".bias"; - bd = ggml_mul_mat(ctx, tensors.at(wn), bd); - bd = ggml_add(ctx, bd, tensors.at(bn)); - if (j < 2) bd = ggml_relu(ctx, bd); + if (!sam3_graph_compute(model.backend, graph, state.n_threads)) { + ggml_gallocr_free(alloc); + ggml_free(ctx); + return result; } - // bd: [4, NQ, 1] - // ref_boxes = sigmoid(inverse_sigmoid(ref_boxes) + box_delta) - auto* ref_inv_cur = sam3_inverse_sigmoid(ctx, ref_boxes); - ref_boxes = ggml_sigmoid(ctx, ggml_add(ctx, ref_inv_cur, bd)); - sam3_name_tensorf(ref_boxes, "ddec_layer%d_refboxes", i); - } + ggml_backend_tensor_get(dout.class_scores, scores_data.data(), 0, NQ * sizeof(float)); + ggml_backend_tensor_get(dout.pred_boxes, boxes_data.data(), 0, 4 * NQ * sizeof(float)); + ggml_backend_tensor_get(dout.presence_score, &presence_logit, 0, sizeof(float)); + ggml_backend_tensor_get(dout.queries, queries_data.data(), 0, + D * (NQ + 1) * sizeof(float)); - // ── Final normalization ────────────────────────────────────────────── - // Match Python: decoder.norm is applied to object queries only. - auto* obj_queries = ggml_view_3d(ctx, queries, D, NQ, 1, - queries->nb[1], queries->nb[2], 1 * queries->nb[1]); - obj_queries = ggml_cont(ctx, obj_queries); - obj_queries = sam3_layer_norm(ctx, obj_queries, - tensors.at("ddec.norm.weight"), - tensors.at("ddec.norm.bias")); - ggml_set_name(obj_queries, "ddec_normed_output"); + ggml_gallocr_free(alloc); + ggml_free(ctx); + } - auto* queries_for_seg = ggml_concat(ctx, last_presence, obj_queries, 1); + float presence_prob = 1.0f / (1.0f + expf(-presence_logit)); - auto* class_scores = sam3_dot_product_scoring(ctx, model, obj_queries, text_feats, text_valid_mask); - // class_scores: [NQ, B] + SAM3_LOG(2, "%s: DETR decoder done\n", __func__); + /* + ** ── SUB-GRAPH 5: Segmentation Head ─────────────────────────────── + */ + const int mask_hw = H * 4; // 288 for SAM3 + std::vector all_masks(NQ * mask_hw * mask_hw); + { + const size_t sz = ggml_tensor_overhead() * 16384 + ggml_graph_overhead() * 2; + struct ggml_init_params gp = {sz, nullptr, true}; + auto* ctx = ggml_init(gp); - // ── Presence score ─────────────────────────────────────────────────── - // Presence token head: LN + 3-layer MLP (D→D→D→1) - auto* pres_out = sam3_layer_norm(ctx, last_presence, - tensors.at("ddec.presence_token_out_norm.weight"), - tensors.at("ddec.presence_token_out_norm.bias")); + auto* enc_h = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, N_spatial, 1); + ggml_set_name(enc_h, "seg_enc"); + ggml_set_input(enc_h); - for (int j = 0; j < 3; ++j) { - auto wn = "ddec.presence_token_head.layers." + std::to_string(j) + ".weight"; - auto bn = "ddec.presence_token_head.layers." + std::to_string(j) + ".bias"; - pres_out = ggml_mul_mat(ctx, tensors.at(wn), pres_out); - pres_out = ggml_add(ctx, pres_out, tensors.at(bn)); - if (j < 2) pres_out = ggml_relu(ctx, pres_out); - } - // Keep presence as raw logit (no sigmoid yet — applied during post-processing) - auto* presence_score = ggml_reshape_2d(ctx, pres_out, 1, 1); - // presence_score: [1, B] — raw logit + // FPN features at 3 scales — create fresh inputs + const int H0 = H * 4, H1 = H * 2; + auto* fpn0 = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, D, H0, H0, 1); + ggml_set_name(fpn0, "seg_fpn0"); + ggml_set_input(fpn0); + auto* fpn1 = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, D, H1, H1, 1); + ggml_set_name(fpn1, "seg_fpn1"); + ggml_set_input(fpn1); + auto* fpn2 = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, D, H, H, 1); + ggml_set_name(fpn2, "seg_fpn2"); + ggml_set_input(fpn2); + struct ggml_tensor* fpn_feats[3] = {fpn0, fpn1, fpn2}; - sam3_ddec_output out; - out.queries = queries_for_seg; // [D, NQ+1, B] - out.presence_feats = last_presence; // [D, 1, B] - out.pred_boxes = ref_boxes; // [4, NQ, B] - out.class_scores = class_scores; // [NQ, B] - out.presence_score = presence_score; // [1, B] + // Object queries (skip presence token at index 0 → start at index 1) + auto* oq = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, NQ, 1); + ggml_set_name(oq, "seg_queries"); + ggml_set_input(oq); - return out; -} + auto* txt = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, T, 1); + ggml_set_name(txt, "seg_text"); + ggml_set_input(txt); + auto* tab = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, T, 1, 1); + ggml_set_name(tab, "seg_bias"); + ggml_set_input(tab); -/***************************************************************************** -** Segmentation head (MaskFormer) — graph building -*****************************************************************************/ + auto* out = sam3_build_seg_head_graph(ctx, model, enc_h, fpn_feats, + oq, txt, tab); + ggml_set_output(out); -// Build pixel decoder: progressively upsample FPN features. -// fpn_feats[0]: [D, 288, 288, B] (highest res) -// fpn_feats[1]: [D, 144, 144, B] -// fpn_feats[2]: [D, 72, 72, B] (lowest res) -// Returns: [D, 288, 288, B] pixel features -// -// Python PixelDecoder.forward: -// prev_fpn = backbone_feats[-1] (lowest res) -// for bb_feat in backbone_feats[:-1][::-1]: (iterate from second-lowest to highest) -// prev_fpn = bb_feat + F.interpolate(prev_fpn, size=bb_feat.shape[-2:], mode="nearest") -// prev_fpn = conv_layers[i](prev_fpn) # conv on the MERGED result -// prev_fpn = F.relu(norms[i](prev_fpn)) # GroupNorm then ReLU -// -// Python uses GroupNorm(8, 256) — we use ggml_group_norm which normalizes ne[2] -// (the channel dim) in groups. The conv output is [W, H, D, B] with D in ne[2]. -static struct ggml_tensor* sam3_pixel_decoder( - struct ggml_context* ctx, - const sam3_model& model, - struct ggml_tensor* fpn_feats[3]) // [D, W, H, B] at 3 scales -{ - const auto& seg = model.seg_head; + auto* graph = ggml_new_graph_custom(ctx, 32768, false); + ggml_build_forward_expand(graph, out); - // Start from lowest resolution - auto* feat = fpn_feats[2]; // [D, 72, 72, B] + SAM3_LOG(2, "%s: seg head graph: %d nodes\n", __func__, ggml_graph_n_nodes(graph)); - // Iteration 0: merge with FPN[1] (144x144) - // prev_fpn = FPN[1] + upsample(prev_fpn) - // Permute to [W, H, D, B] for conv operations - auto* prev = ggml_cont(ctx, ggml_permute(ctx, feat, 2, 0, 1, 3)); // [72, 72, D, B] - prev = ggml_upscale(ctx, prev, 2, GGML_SCALE_MODE_NEAREST); // [144, 144, D, B] - auto* fpn1 = ggml_cont(ctx, ggml_permute(ctx, fpn_feats[1], 2, 0, 1, 3)); // [144, 144, D, B] - prev = ggml_add(ctx, fpn1, prev); // merged - // Conv 3x3 on the MERGED result (not individual FPN feat) - prev = ggml_conv_2d_s1_ph(ctx, seg.up_conv_w[0], prev); - { - auto* b3d = ggml_reshape_3d(ctx, seg.up_conv_b[0], 1, 1, seg.up_conv_b[0]->ne[0]); - prev = ggml_add(ctx, prev, ggml_repeat(ctx, b3d, prev)); - } - // GroupNorm(8, 256) then ReLU — prev is [W, H, D, B] with D in ne[2] - prev = ggml_group_norm(ctx, prev, 8, 1e-5f); - { - auto* w3d = ggml_reshape_3d(ctx, seg.up_norm_w[0], 1, 1, seg.up_norm_w[0]->ne[0]); - prev = ggml_mul(ctx, prev, ggml_repeat(ctx, w3d, prev)); - auto* bn3d = ggml_reshape_3d(ctx, seg.up_norm_b[0], 1, 1, seg.up_norm_b[0]->ne[0]); - prev = ggml_add(ctx, prev, ggml_repeat(ctx, bn3d, prev)); - } - prev = ggml_relu(ctx, prev); - ggml_set_name(prev, "seg_pixel_dec_stage0"); + auto* alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); + if (!ggml_gallocr_reserve(alloc, graph) || !ggml_gallocr_alloc_graph(alloc, graph)) { + fprintf(stderr, "%s: segmentation head alloc failed\n", __func__); + ggml_gallocr_free(alloc); + ggml_free(ctx); + return result; + } - // Iteration 1: merge with FPN[0] (288x288) - prev = ggml_upscale(ctx, prev, 2, GGML_SCALE_MODE_NEAREST); // [288, 288, D, B] - auto* fpn0 = ggml_cont(ctx, ggml_permute(ctx, fpn_feats[0], 2, 0, 1, 3)); // [288, 288, D, B] - prev = ggml_add(ctx, fpn0, prev); // merged - // Conv 3x3 on the MERGED result - prev = ggml_conv_2d_s1_ph(ctx, seg.up_conv_w[1], prev); - { - auto* b3d = ggml_reshape_3d(ctx, seg.up_conv_b[1], 1, 1, seg.up_conv_b[1]->ne[0]); - prev = ggml_add(ctx, prev, ggml_repeat(ctx, b3d, prev)); - } - // GroupNorm(8, 256) then ReLU - prev = ggml_group_norm(ctx, prev, 8, 1e-5f); - { - auto* w3d = ggml_reshape_3d(ctx, seg.up_norm_w[1], 1, 1, seg.up_norm_w[1]->ne[0]); - prev = ggml_mul(ctx, prev, ggml_repeat(ctx, w3d, prev)); - auto* bn3d = ggml_reshape_3d(ctx, seg.up_norm_b[1], 1, 1, seg.up_norm_b[1]->ne[0]); - prev = ggml_add(ctx, prev, ggml_repeat(ctx, bn3d, prev)); - } - prev = ggml_relu(ctx, prev); - ggml_set_name(prev, "seg_pixel_dec_stage1"); + ggml_backend_tensor_set(enc_h, fenc_output_cpu.data(), 0, + D * N_spatial * sizeof(float)); - // Python PixelDecoder allocates 3 conv layers but only uses 2 (one per - // upsample step). The 3rd conv (up_conv_w[2]) is unused. + // Copy FPN features from state + { + size_t s0 = (size_t)D * H0 * H0 * sizeof(float); + size_t s1 = (size_t)D * H1 * H1 * sizeof(float); + size_t s2 = (size_t)D * H * H * sizeof(float); + std::vector b0(D * H0 * H0), b1(D * H1 * H1), b2(D * H * H); + ggml_backend_tensor_get(state.neck_det[0], b0.data(), 0, s0); + if (fpn0->buffer) ggml_backend_tensor_set(fpn0, b0.data(), 0, s0); + ggml_backend_tensor_get(state.neck_det[1], b1.data(), 0, s1); + if (fpn1->buffer) ggml_backend_tensor_set(fpn1, b1.data(), 0, s1); + ggml_backend_tensor_get(state.neck_det[2], b2.data(), 0, s2); + if (fpn2->buffer) ggml_backend_tensor_set(fpn2, b2.data(), 0, s2); + } - auto* out = ggml_cont(ctx, ggml_permute(ctx, prev, 1, 2, 0, 3)); // [D, 288, 288, B] - return out; -} + // Object queries: extract from DETR queries (skip presence token at slot 0) + // queries_data is flat [D * 201], presence token is at positions [0..D-1] + // Object queries start at position [D..D*(NQ+1)-1] + ggml_backend_tensor_set(oq, queries_data.data() + D, 0, D * NQ * sizeof(float)); -// Build the full segmentation head graph. -// -// Python UniversalSegmentationHead.forward: -// 1. Cross-attend encoder_hidden_states to prompt → updated encoder -// 2. _embed_pixels: replace lowest-res FPN feat with spatial portion of encoder output -// 3. Run pixel decoder on modified FPN feats -// 4. instance_seg_head (Conv1x1) -// 5. mask_predictor: einsum(mask_embed(queries), instance_embeds) -// -// enc_hidden: [D, N_spatial, B] — fusion encoder output (cross-attended in step 1) -// fpn_feats[3]: the 3 FPN features at different resolutions -// query_outputs: [D, N, B] selected object query outputs -// text_features: [D, T, B] for cross-attention (prompt) -// Returns: mask_logits [W*H, N, B] (raw logits, not sigmoid) -static struct ggml_tensor* sam3_build_seg_head_graph( - struct ggml_context* ctx, - const sam3_model& model, - struct ggml_tensor* enc_hidden, // [D, N_spatial, B] fusion encoder output - struct ggml_tensor* fpn_feats[3], // FPN features at 3 scales - struct ggml_tensor* query_outputs, // [D, N, B] - struct ggml_tensor* text_features, // [D, T, B] (for cross-attn, can be nullptr) - struct ggml_tensor* text_attn_bias = nullptr) { - const auto& seg = model.seg_head; - const auto& tensors = model.tensors; - const int64_t D = enc_hidden->ne[0]; // 256 - const int64_t B = enc_hidden->ne[2]; // 1 - const int64_t N = query_outputs->ne[1]; // number of selected queries + ggml_backend_tensor_set(txt, combined_prompt_cpu.data(), 0, D * T * sizeof(float)); + ggml_backend_tensor_set(tab, combined_bias_cpu.data(), 0, T * sizeof(float)); - auto* enc = enc_hidden; - if (text_features) { - auto* ca_norm = sam3_layer_norm(ctx, enc, - tensors.at("seg.cross_attn_norm.weight"), - tensors.at("seg.cross_attn_norm.bias")); + if (!sam3_graph_compute(model.backend, graph, state.n_threads)) { + ggml_gallocr_free(alloc); + ggml_free(ctx); + return result; + } + ggml_backend_tensor_get(out, all_masks.data(), 0, all_masks.size() * sizeof(float)); - auto* ca_mask = sam3_expand_token_attn_bias(ctx, text_attn_bias, enc->ne[1], 8, B); - auto* ca_out = sam3_multihead_attn_fused(ctx, ca_norm, text_features, - seg.ca_prompt_q_w, seg.ca_prompt_q_b, - seg.ca_prompt_out_w, seg.ca_prompt_out_b, - 8, ca_mask); - enc = ggml_add(ctx, enc, ca_out); + ggml_gallocr_free(alloc); + ggml_free(ctx); } - // enc: [D, N_spatial, B] - ggml_set_name(enc, "seg_enc_after_ca"); - // Replace lowest-res FPN feat with spatial portion of encoder output - const int64_t feat_hw = model.hparams.n_img_embd(); // 72 - auto* enc_spatial = ggml_reshape_4d(ctx, enc, D, feat_hw, feat_hw, B); -#ifndef NDEBUG - auto* enc_spatial_dbg = ggml_cont(ctx, ggml_permute(ctx, enc_spatial, 2, 0, 1, 3)); - ggml_set_name(enc_spatial_dbg, "seg_enc_visual"); -#endif + /* + ** ── Post-processing: thresholding + NMS + mask resize ──────────── + */ + std::vector dets; + for (int q = 0; q < NQ; ++q) { + float class_prob = 1.0f / (1.0f + expf(-scores_data[q])); + float score = class_prob * presence_prob; + if (score < params.score_threshold) continue; - struct ggml_tensor* modified_fpn[3] = { - fpn_feats[0], - fpn_feats[1], - enc_spatial, // replaces original lowest-res FPN - }; + sam3_detection det; + float cx = boxes_data[0 + q * 4]; + float cy = boxes_data[1 + q * 4]; + float bw = boxes_data[2 + q * 4]; + float bh = boxes_data[3 + q * 4]; - auto* pixel_feats = sam3_pixel_decoder(ctx, model, modified_fpn); -#ifndef NDEBUG - auto* pixel_feats_dbg = ggml_cont(ctx, ggml_permute(ctx, pixel_feats, 2, 0, 1, 3)); - ggml_set_name(pixel_feats_dbg, "seg_pixel_decoder_out"); -#endif + det.box = sam3_cxcywh_to_xyxy(cx, cy, bw, bh, state.orig_width, state.orig_height); + det.score = score; - const int64_t W = pixel_feats->ne[1]; // 288 - const int64_t H = pixel_feats->ne[2]; // 288 + const float* mask_ptr = all_masks.data() + q * mask_hw * mask_hw; + auto mask_resized = sam3_bilinear_interpolate(mask_ptr, mask_hw, mask_hw, + state.orig_width, state.orig_height); + det.mask.width = state.orig_width; + det.mask.height = state.orig_height; + det.mask.data.resize(state.orig_width * state.orig_height); + for (int i = 0; i < (int)mask_resized.size(); ++i) + det.mask.data[i] = (mask_resized[i] > 0.0f) ? 255 : 0; + det.mask.iou_score = score; - // Instance segmentation head (Conv1x1) - auto* pf_conv = ggml_cont(ctx, ggml_permute(ctx, pixel_feats, 2, 0, 1, 3)); - pf_conv = ggml_conv_2d_sk_p0(ctx, tensors.at("seg.instance_seg_head.weight"), pf_conv); - { - auto* b3d = ggml_reshape_3d(ctx, tensors.at("seg.instance_seg_head.bias"), - 1, 1, tensors.at("seg.instance_seg_head.bias")->ne[0]); - pf_conv = ggml_add(ctx, pf_conv, ggml_repeat(ctx, b3d, pf_conv)); + dets.push_back(std::move(det)); } - auto* pixel_embed = ggml_cont(ctx, ggml_permute(ctx, pf_conv, 1, 2, 0, 3)); // [D, W, H, B] -#ifndef NDEBUG - auto* pixel_embed_dbg = ggml_cont(ctx, ggml_permute(ctx, pixel_embed, 2, 0, 1, 3)); - ggml_set_name(pixel_embed_dbg, "seg_instance_embed"); -#endif - // Mask embedding MLP - auto* mask_embed = query_outputs; - for (int j = 0; j < 3; ++j) { - auto wn = "seg.mask_predictor.mask_embed.layers." + std::to_string(j) + ".weight"; - auto bn = "seg.mask_predictor.mask_embed.layers." + std::to_string(j) + ".bias"; - mask_embed = ggml_mul_mat(ctx, tensors.at(wn), mask_embed); - mask_embed = ggml_add(ctx, mask_embed, tensors.at(bn)); - if (j < 2) mask_embed = ggml_relu(ctx, mask_embed); + SAM3_LOG(2, "%s: %zu detections above threshold %.2f (presence=%.3f, logit=%.3f)\n", + __func__, dets.size(), params.score_threshold, presence_prob, presence_logit); + + auto keep = sam3_nms(dets, params.nms_threshold); + for (int i = 0; i < (int)keep.size(); ++i) { + dets[keep[i]].instance_id = i + 1; + result.detections.push_back(std::move(dets[keep[i]])); } - ggml_set_name(mask_embed, "seg_mask_embed"); - // Mask prediction: einsum('bqc,bchw->bqhw') - auto* pe_flat = ggml_reshape_3d(ctx, pixel_embed, D, W * H, B); - auto* masks = ggml_mul_mat(ctx, pe_flat, mask_embed); - ggml_set_name(masks, "seg_mask_logits"); + SAM3_LOG(2, "%s: %zu detections after NMS\n", __func__, result.detections.size()); - return masks; +#if SAM3_LOG_LEVEL >= 1 + auto t_end = std::chrono::high_resolution_clock::now(); + double total_ms = std::chrono::duration(t_end - t_start).count(); + SAM3_LOG(1, "%s: completed in %.1f ms\n", __func__, total_ms); +#endif + + return result; } /***************************************************************************** -** Memory encoder (Phase 7, Step 7.1) +** SAM attention helper (separate Q, K, V weight/bias) *****************************************************************************/ -// CXBlock: depthwise conv + LayerNorm + pointwise MLP with residual scaling. -static struct ggml_tensor* sam3_cxblock_forward( +static struct ggml_tensor* sam3_sam_attention( struct ggml_context* ctx, - struct ggml_tensor* x, // [D, H, W, B] - struct ggml_tensor* dw_w, // [7, 7, 1, D] depthwise - struct ggml_tensor* dw_b, // [D] - struct ggml_tensor* norm_w, // [D] - struct ggml_tensor* norm_b, // [D] - struct ggml_tensor* fc1_w, // [D, 1024] - struct ggml_tensor* fc1_b, // [1024] - struct ggml_tensor* fc2_w, // [1024, D] - struct ggml_tensor* fc2_b, // [D] - struct ggml_tensor* gamma) // [D] -{ - const int D = (int)x->ne[0]; - const int H = (int)x->ne[1]; - const int W = (int)x->ne[2]; - - // ggml conv expects WHCB layout; internal feature maps are stored as CWHB. - auto* x_whcb = ggml_cont(ctx, ggml_permute(ctx, x, 2, 0, 1, 3)); - - // Depthwise conv (groups = D): use the direct depthwise path. - // ggml_conv_2d_dw_direct only supports f32 kernel — cast if needed. - auto* dw_w_f32 = (dw_w->type == GGML_TYPE_F32) ? dw_w : ggml_cast(ctx, dw_w, GGML_TYPE_F32); - auto* h = ggml_conv_2d_dw_direct(ctx, dw_w_f32, x_whcb, 1, 1, 3, 3, 1, 1); - h = ggml_add(ctx, h, ggml_reshape_4d(ctx, dw_b, 1, 1, D, 1)); - h = ggml_cont(ctx, ggml_permute(ctx, h, 1, 2, 0, 3)); + struct ggml_tensor* q_in, // [D, N_q, B] + struct ggml_tensor* k_in, // [D, N_kv, B] + struct ggml_tensor* v_in, // [D, N_kv, B] + const sam3_sam_attn& attn, + int n_heads) { + const int64_t N_q = q_in->ne[1]; + const int64_t B = q_in->ne[2]; + const int64_t N_kv = k_in->ne[1]; - // LayerNorm2d - h = sam3_layer_norm_2d(ctx, h, norm_w, norm_b); + // Project + auto* Q = ggml_add(ctx, ggml_mul_mat(ctx, attn.q_w, q_in), attn.q_b); + auto* K = ggml_add(ctx, ggml_mul_mat(ctx, attn.k_w, k_in), attn.k_b); + auto* V = ggml_add(ctx, ggml_mul_mat(ctx, attn.v_w, v_in), attn.v_b); - // Pointwise MLP: reshape to [D, H*W, B], apply FC, reshape back - auto* flat = ggml_reshape_3d(ctx, h, D, H * W, 1); - flat = ggml_add(ctx, ggml_mul_mat(ctx, fc1_w, flat), fc1_b); - flat = ggml_gelu(ctx, flat); - flat = ggml_add(ctx, ggml_mul_mat(ctx, fc2_w, flat), fc2_b); - h = ggml_reshape_4d(ctx, flat, D, H, W, 1); + // Debug: mark projections for the first SA call (N_q=8 tokens, block 0) + static int _sa_call_count = 0; + if (_sa_call_count == 0 && N_q <= 16) { + ggml_set_name(Q, "dbg_sa0_Q_proj"); + ggml_set_output(Q); + ggml_set_name(V, "dbg_sa0_V_proj"); + ggml_set_output(V); + } + _sa_call_count++; - // Residual with learnable scaling: x + gamma * h - auto* gamma_4d = ggml_reshape_4d(ctx, gamma, D, 1, 1, 1); - h = ggml_mul(ctx, h, gamma_4d); - return ggml_add(ctx, x, h); -} + // internal_dim = out_proj cols = attn.q_w->ne[1] + const int64_t ID = attn.q_w->ne[1]; + const int64_t HD = ID / n_heads; -/***************************************************************************** -** Prompt building for memory attention (Phase 7) -*****************************************************************************/ + // Reshape to multi-head: [HD, N, NH, B] + Q = ggml_reshape_4d(ctx, Q, HD, n_heads, N_q, B); + Q = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 1, 3)); // [HD, N_q, NH, B] -struct sam3_prompt_data { - std::vector prompt; // [MD * M_total] flat - std::vector prompt_pos; // [MD * M_total] flat - int M_total; // total prompt tokens - int num_obj_ptr_tokens; // pointer tokens at end (excluded from RoPE) - int M_spatial; // spatial memory tokens = M_total - num_obj_ptr_tokens -}; + K = ggml_reshape_4d(ctx, K, HD, n_heads, N_kv, B); + K = ggml_cont(ctx, ggml_permute(ctx, K, 0, 2, 1, 3)); // [HD, N_kv, NH, B] -// Build prompt and prompt_pos tensors CPU-side for memory attention. -// mem_slot_feats[i]: [MD * H * H] raw spatial features (ggml layout) -// mem_slot_pes[i]: [MD * H * H] sinusoidal spatial PE (ggml layout) -// spatial_tpos[i]: temporal position for slot i (0 = conditioning, >0 = non-cond t_pos) -// obj_ptrs[i]: [D] object pointer data -// ptr_tpos[i]: temporal position for pointer i (<0 = skip) -static sam3_prompt_data sam3_build_prompt_and_pos( - const sam3_model& model, - const std::vector>& mem_slot_feats, - const std::vector>& mem_slot_pes, - const std::vector& spatial_tpos, - const std::vector>& obj_ptrs, - const std::vector& ptr_tpos, - int eff_feat_size = 0) { - const auto& hp = model.hparams; - const int MD = hp.mem_out_dim; // 64 - const int D = hp.neck_dim; // 256 - const int H = (eff_feat_size > 0) ? eff_feat_size : hp.feat_size(); - // Per-slot token count: from actual data size (handles both HxH and perceiver 512) - const int HH = mem_slot_feats.empty() ? (H * H) - : (int)(mem_slot_feats[0].size() / MD); + V = ggml_reshape_4d(ctx, V, HD, n_heads, N_kv, B); + V = ggml_cont(ctx, ggml_permute(ctx, V, 0, 2, 1, 3)); // [HD, N_kv, NH, B] contiguous - sam3_prompt_data pd; - pd.num_obj_ptr_tokens = 0; + // Attention + float scale = 1.0f / sqrtf((float)HD); + auto* out = ggml_flash_attn_ext(ctx, Q, K, V, nullptr, scale, 0.0f, 0.0f); + // out: [HD, NH, N_q, B] (flash_attn_ext swaps dims 1,2 vs input) - // Read maskmem_tpos_enc from model (one tensor [MD, 1, 1, 7]) - std::vector tpos_all(MD * hp.num_maskmem); - if (model.mem_enc.tpos[0]) { - sam3_read_f32(model.mem_enc.tpos[0], tpos_all.data(), MD * hp.num_maskmem); - } +#if 0 // Manual SDPA (for debugging only) + auto* Q3 = ggml_reshape_3d(ctx, Q, HD, N_q, n_heads * B); + auto* K3 = ggml_reshape_3d(ctx, K, HD, N_kv, n_heads * B); + auto* V3 = ggml_reshape_3d(ctx, V, HD, N_kv, n_heads * B); + // QK^T: ggml_mul_mat(K, Q) → K^T @ Q → [N_kv, N_q, NH*B] + auto* attn_scores = ggml_mul_mat(ctx, K3, Q3); + attn_scores = ggml_scale(ctx, attn_scores, scale); + attn_scores = ggml_soft_max(ctx, attn_scores); - // 1. Spatial memory tokens - for (size_t s = 0; s < mem_slot_feats.size(); ++s) { - const int slot_tokens = (int)(mem_slot_feats[s].size() / MD); - // Append spatial features - pd.prompt.insert(pd.prompt.end(), mem_slot_feats[s].begin(), mem_slot_feats[s].end()); + // attn @ V: need attn^T [N_q, N_kv] and V^T [HD, N_kv] + // ggml_mul_mat(attn^T, V) = (attn^T)^T @ V = attn @ V = [N_q, HD]... no. + // ggml_mul_mat(A, B) = A^T @ B where A=[K, M], B=[K, N] → [M, N] + // Want: output[q, d] = sum_k attn[q, k] * V[k, d] + // = (V^T @ attn^T)^T... let me think differently. + // attn_scores is [N_kv, N_q, NH*B]. For each head: + // attn[k, q] = attn_scores[k, q] (col q has the weights for query q) + // V3 is [HD, N_kv, NH*B]. + // Want: out[d, q] = sum_k V[d, k] * attn[k, q] = V @ attn + // = ggml_mul_mat? mul_mat(A, B) = A^T B with A=[K, M], B=[K, N] → [M, N] + // V has ne=[HD, N_kv, ...]. attn has ne=[N_kv, N_q, ...]. + // If A=V3 (ne0=HD, ne1=N_kv) and B=attn_scores (ne0=N_kv, ne1=N_q): + // Shared dim ne0: V3 ne0=HD ≠ attn ne0=N_kv. Mismatch! + // + // Need to transpose V: V^T is [N_kv, HD]. Then A=V^T, B=attn_scores. + // A ne0=N_kv, B ne0=N_kv → shared. A^T B = V @ attn → [HD, N_q]. ✓ + auto* VT = ggml_permute(ctx, V3, 1, 0, 2, 3); // [N_kv, HD, NH*B] + VT = ggml_cont(ctx, VT); + auto* out3 = ggml_mul_mat(ctx, VT, attn_scores); // [HD, N_q, NH*B] - // Build positional encoding: spatial PE + temporal PE - std::vector pos = mem_slot_pes[s]; // copy spatial PE - int tpos_idx = spatial_tpos[s]; - if (tpos_idx >= 0 && tpos_idx < hp.num_maskmem) { - int enc_idx = hp.num_maskmem - tpos_idx - 1; // Python: maskmem_tpos_enc[num_maskmem - tpos - 1] - for (int n = 0; n < slot_tokens; ++n) - for (int d = 0; d < MD; ++d) - pos[d + n * MD] += tpos_all[d + enc_idx * MD]; - } - pd.prompt_pos.insert(pd.prompt_pos.end(), pos.begin(), pos.end()); - } - pd.M_spatial = (int)mem_slot_feats.size() * HH; + // Reshape back to 4D: [HD, N_q, NH, B] + auto* out = ggml_reshape_4d(ctx, out3, HD, N_q, n_heads, B); + // Permute to [HD, NH, N_q, B] to match flash_attn_ext output convention + out = ggml_cont(ctx, ggml_permute(ctx, out, 0, 2, 1, 3)); +#endif - // 2. Object pointer tokens (split each D=256 pointer into D/MD=4 tokens of MD=64) - const int split = D / MD; // 4 + // Merge heads: [ID=HD*NH, N_q, B] + auto* merged = ggml_reshape_3d(ctx, out, ID, N_q, B); - // Read obj_ptr_tpos_proj weights for CPU-side matmul - std::vector tpos_w(D * MD), tpos_b(MD); - if (model.obj_ptr_tpos_w) { - sam3_read_f32(model.obj_ptr_tpos_w, tpos_w.data(), D * MD); - sam3_read_f32(model.obj_ptr_tpos_b, tpos_b.data(), MD); + // Debug: mark merged attention output for first SA call + static int _sa_merge_count = 0; + if (_sa_merge_count == 0 && N_q <= 16) { + ggml_set_name(merged, "dbg_sa0_merged"); + ggml_set_output(merged); } + _sa_merge_count++; - for (size_t p = 0; p < obj_ptrs.size(); ++p) { - int rel = ptr_tpos[p]; - if (rel < 0) continue; - - // 1D sine PE for temporal position (normalized by max_obj_ptrs - 1 = 15) - std::vector sine_pe(D); - sam3_get_1d_sine_pe(sine_pe.data(), (float)rel / 15.0f, D); - - // Project 256-dim sine PE → 64-dim via obj_ptr_tpos_proj (CPU matmul) - // W is [D=256, MD=64] in ggml; y[j] = sum_i W[i + j*D] * x[i] + b[j] - std::vector proj_pe(MD, 0.0f); - for (int j = 0; j < MD; ++j) { - float sum = tpos_b[j]; - for (int i = 0; i < D; ++i) - sum += tpos_w[i + j * D] * sine_pe[i]; - proj_pe[j] = sum; - } - - // Split pointer [256] → 4 tokens of [64] - // Python: ptr.view(1,1,4,64).permute(1,2,0,3).reshape(4,1,64) - const auto& ptr = obj_ptrs[p]; - for (int t = 0; t < split; ++t) { - for (int d = 0; d < MD; ++d) { - pd.prompt.push_back(ptr[t * MD + d]); - pd.prompt_pos.push_back(proj_pe[d]); - } - pd.num_obj_ptr_tokens++; - } - } + // Output projection + out = ggml_mul_mat(ctx, attn.out_w, merged); + out = ggml_add(ctx, out, attn.out_b); - pd.M_total = pd.M_spatial + pd.num_obj_ptr_tokens; - return pd; + return out; } /***************************************************************************** -** Memory attention (Phase 7, Step 7.2) +** SAM prompt encoder — graph building (Phase 6, Step 6.1) *****************************************************************************/ -// Build memory attention graph. -// curr_tokens: [D, N, 1] — current frame image tokens (flattened from 72x72 = 5184) -// src_pos: [D, N, 1] — sinusoidal PE for current tokens (pos_enc_at_input) -// prompt: [MD, M_total, 1] — concatenated memory features + split pointer tokens -// prompt_pos: [MD, M_total, 1] — positional encodings for prompt -// rope_freqs: [2, D/2, N] — axial RoPE frequencies for Q and self-attn K -// rope_k_freqs:[2, D/2, M_spatial] — repeated axial RoPE for cross-attn spatial K -// num_obj_ptr_tokens: number of pointer tokens at end of prompt (excluded from RoPE) -// Returns: conditioned tokens [D, N, 1] -static struct ggml_tensor* sam3_build_mem_attn_graph( - struct ggml_context* ctx, - const sam3_model& model, - struct ggml_tensor* curr_tokens, // [D, N, 1] - struct ggml_tensor* src_pos, // [D, N, 1] - struct ggml_tensor* prompt, // [MD, M_total, 1] - struct ggml_tensor* prompt_pos, // [MD, M_total, 1] - struct ggml_tensor* rope_freqs, // [2, D/2, N] - struct ggml_tensor* rope_k_freqs, // [2, D/2, M_spatial] or nullptr - int num_obj_ptr_tokens) { - const auto& ma = model.mem_attn; - const int D = model.hparams.neck_dim; // 256 - const int N = (int)curr_tokens->ne[1]; - const int M_total = (int)prompt->ne[1]; - const int M_spatial = M_total - num_obj_ptr_tokens; - - // pos_enc_at_input: x = curr + 0.1 * src_pos - auto* x = ggml_add(ctx, curr_tokens, ggml_scale(ctx, src_pos, 0.1f)); - ggml_set_name(x, "phase7_mem_attn_input"); - - for (int l = 0; l < (int)ma.layers.size(); ++l) { - const auto& ly = ma.layers[l]; - - // ── Self-attention with RoPE ────────────────────────────────────── - { - auto* x_norm = sam3_layer_norm(ctx, x, ly.norm1_w, ly.norm1_b); - auto* q = ggml_add(ctx, ggml_mul_mat(ctx, ly.sa_q_w, x_norm), ly.sa_q_b); - auto* k = ggml_add(ctx, ggml_mul_mat(ctx, ly.sa_k_w, x_norm), ly.sa_k_b); - auto* v = ggml_add(ctx, ggml_mul_mat(ctx, ly.sa_v_w, x_norm), ly.sa_v_b); - - // Apply RoPE to Q and K (single-head, [D, N, 1]) - q = sam3_apply_rope(ctx, q, rope_freqs); - k = sam3_apply_rope(ctx, k, rope_freqs); +// Random Fourier positional encoding for a single (x, y) coordinate +// coords_norm: normalized to [0, 1], pe_gaussian: PyTorch row-major [2, 128] +// Output: [256] = [sin(128); cos(128)] +static void sam3_pe_encode_coord(float* out, float x_norm, float y_norm, + const float* pe_gauss, int num_pos_feats) { + // Map [0,1] → [-1,1] + float coords[2] = {2.0f * x_norm - 1.0f, 2.0f * y_norm - 1.0f}; - // Single-head attention (256-dim) - q = ggml_reshape_4d(ctx, q, D, 1, N, 1); - k = ggml_reshape_4d(ctx, k, D, 1, N, 1); - v = ggml_reshape_4d(ctx, v, D, 1, N, 1); - q = ggml_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3)); - k = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3)); - v = ggml_permute(ctx, v, 0, 2, 1, 3); + // coords @ pe_gaussian → [128] + for (int i = 0; i < num_pos_feats; ++i) { + float dot = coords[0] * pe_gauss[i] + + coords[1] * pe_gauss[num_pos_feats + i]; + dot *= 2.0f * (float)M_PI; + out[i] = sinf(dot); + out[i + num_pos_feats] = cosf(dot); + } +} - float scale = 1.0f / sqrtf((float)D); - auto* sa_out = ggml_flash_attn_ext(ctx, q, k, v, nullptr, scale, 0.0f, 0.0f); - sa_out = ggml_reshape_3d(ctx, sa_out, D, N, 1); - sa_out = ggml_add(ctx, ggml_mul_mat(ctx, ly.sa_out_w, sa_out), ly.sa_out_b); - x = ggml_add(ctx, x, sa_out); - sam3_name_tensorf(x, "phase7_mem_attn_layer%d_after_sa", l); - } +// Read SAM prompt encoder weights from GPU and cache them in state. +// Also pre-computes the dense PE grid and no-mask tiled embedding. +// These never change between PVS calls for the same model. +static void sam3_populate_pe_cache(sam3_state& state, const sam3_model& model) { + if (state.pe_cache_valid) return; - // ── Cross-attention to memory (kv_dim=64→256) with RoPE ────────── - { - auto* x_norm = sam3_layer_norm(ctx, x, ly.norm2_w, ly.norm2_b); - auto* q = ggml_add(ctx, ggml_mul_mat(ctx, ly.ca_q_w, x_norm), ly.ca_q_b); + const int D = model.hparams.sam_embed_dim; // 256 + const int H = sam3_eff_feat_size(state, model.hparams); + const int num_pos_feats = D / 2; // 128 + const int pe_nel = 2 * num_pos_feats; // 256 + const auto& pe = model.sam_pe; - // K: project from (prompt + prompt_pos) — pos enc added before projection - auto* kv_with_pos = ggml_add(ctx, prompt, prompt_pos); - auto* k = ggml_add(ctx, ggml_mul_mat(ctx, ly.ca_k_w, kv_with_pos), ly.ca_k_b); - // V: project from prompt only (no pos enc) - auto* v = ggml_add(ctx, ggml_mul_mat(ctx, ly.ca_v_w, prompt), ly.ca_v_b); - - // Apply RoPE to Q - q = sam3_apply_rope(ctx, q, rope_freqs); - - // Apply RoPE to spatial K only (exclude num_obj_ptr_tokens tail keys) - if (M_spatial > 0 && rope_k_freqs) { - auto* k_spatial = ggml_cont(ctx, ggml_view_3d(ctx, k, - D, M_spatial, 1, k->nb[1], k->nb[2], 0)); - k_spatial = sam3_apply_rope(ctx, k_spatial, rope_k_freqs); - if (num_obj_ptr_tokens > 0) { - auto* k_ptr = ggml_cont(ctx, ggml_view_3d(ctx, k, - D, num_obj_ptr_tokens, 1, k->nb[1], k->nb[2], - (size_t)M_spatial * k->nb[1])); - k = ggml_concat(ctx, k_spatial, k_ptr, 1); - } else { - k = k_spatial; - } - } - - // Flash attention (single head) - q = ggml_reshape_4d(ctx, q, D, 1, N, 1); - k = ggml_reshape_4d(ctx, k, D, 1, M_total, 1); - v = ggml_reshape_4d(ctx, v, D, 1, M_total, 1); - q = ggml_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3)); - k = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3)); - v = ggml_permute(ctx, v, 0, 2, 1, 3); - - float scale = 1.0f / sqrtf((float)D); - auto* ca_out = ggml_flash_attn_ext(ctx, q, k, v, nullptr, scale, 0.0f, 0.0f); - ca_out = ggml_reshape_3d(ctx, ca_out, D, N, 1); - ca_out = ggml_add(ctx, ggml_mul_mat(ctx, ly.ca_out_w, ca_out), ly.ca_out_b); - x = ggml_add(ctx, x, ca_out); - sam3_name_tensorf(x, "phase7_mem_attn_layer%d_after_ca", l); - } - - // ── FFN ─────────────────────────────────────────────────────────── - { - auto* x_norm = sam3_layer_norm(ctx, x, ly.norm3_w, ly.norm3_b); - auto* ffn = ggml_add(ctx, ggml_mul_mat(ctx, ly.ffn_fc1_w, x_norm), ly.ffn_fc1_b); - ffn = ggml_relu(ctx, ffn); - ffn = ggml_add(ctx, ggml_mul_mat(ctx, ly.ffn_fc2_w, ffn), ly.ffn_fc2_b); - x = ggml_add(ctx, x, ffn); - sam3_name_tensorf(x, "phase7_mem_attn_layer%d_after_ffn", l); - } - } - - // Final norm - auto* norm_w = model.tensors.at("mem_attn.norm.weight"); - auto* norm_b = model.tensors.at("mem_attn.norm.bias"); - x = sam3_layer_norm(ctx, x, norm_w, norm_b); - ggml_set_name(x, "phase7_mem_attn_output"); - - return x; // [D, N, 1] -} - -/***************************************************************************** -** Object pointer extraction (Phase 7, Step 7.3) -*****************************************************************************/ - -// Extract object pointer from SAM output token via 3-layer MLP (CPU-side). -static void sam3_extract_obj_ptr_cpu( - const sam3_model& model, - const float* sam_token_data, // [D] - float obj_score, - float* out_ptr) // [D] -{ - const auto& hp = model.hparams; - const int D = hp.neck_dim; - - // SAM2/EdgeTAM with fixed_no_obj_ptr: blend projected ptr with no_obj_ptr - // based on presence score λ. - // SAM3 / SAM2 without fixed_no_obj_ptr: binary threshold. - if ((hp.is_sam2() || hp.is_edgetam()) && hp.fixed_no_obj_ptr) { - // λ = (obj_score > 0) ? 1.0 : 0.0 (hard threshold, not sigmoid) - float lambda = (obj_score > 0.0f) ? 1.0f : 0.0f; - - // Project token through MLP - std::vector h(D), tmp(D); - std::copy(sam_token_data, sam_token_data + D, h.data()); - for (int j = 0; j < 3; ++j) { - auto* w = model.obj_ptr_proj_w[j]; - auto* b = model.obj_ptr_proj_b[j]; - int nel_w = (int)(w->ne[0] * w->ne[1]); - std::vector w_data(nel_w), b_data(D); - sam3_read_f32(w, w_data.data(), nel_w); - sam3_read_f32(b, b_data.data(), D); - for (int o = 0; o < D; ++o) { - float sum = b_data[o]; - for (int i = 0; i < D; ++i) sum += w_data[o * D + i] * h[i]; - tmp[o] = (j < 2) ? std::max(0.0f, sum) : sum; - } - std::swap(h, tmp); - } - - // Blend: obj_ptr = λ * projected + (1-λ) * no_obj_ptr - std::vector no_ptr(D); - ggml_backend_tensor_get(model.no_obj_ptr, no_ptr.data(), 0, D * sizeof(float)); - for (int i = 0; i < D; ++i) - out_ptr[i] = lambda * h[i] + (1.0f - lambda) * no_ptr[i]; - return; - } - - // SAM3 / default path: binary threshold - if (obj_score <= 0.0f) { - ggml_backend_tensor_get(model.no_obj_ptr, out_ptr, 0, D * sizeof(float)); - return; - } - - std::vector h(D), tmp(D); - std::copy(sam_token_data, sam_token_data + D, h.data()); - - for (int j = 0; j < 3; ++j) { - auto* w = model.obj_ptr_proj_w[j]; - auto* b = model.obj_ptr_proj_b[j]; - - int nel_w = (int)(w->ne[0] * w->ne[1]); - std::vector w_data(nel_w); - sam3_read_f32(w, w_data.data(), nel_w); - - std::vector b_data(D); - sam3_read_f32(b, b_data.data(), D); - - for (int o = 0; o < D; ++o) { - float sum = b_data[o]; - for (int i = 0; i < D; ++i) { - sum += w_data[o * D + i] * h[i]; - } - tmp[o] = (j < 2) ? std::max(0.0f, sum) : sum; - } - std::swap(h, tmp); - } - std::copy(h.begin(), h.end(), out_ptr); -} - -/***************************************************************************** -** Tracker infrastructure (Phase 7, Step 7.4) -*****************************************************************************/ - -// Select memory frames for propagation (most recent + evenly spaced). -static std::vector sam3_select_memory_frames( - const std::vector& bank, - int max_slots) { - if ((int)bank.size() <= max_slots) { - std::vector all(bank.size()); - for (int i = 0; i < (int)bank.size(); ++i) all[i] = i; - return all; - } - std::vector selected; - selected.push_back(0); - selected.push_back((int)bank.size() - 1); - int remaining = max_slots - 2; - if (remaining > 0) { - float step = (float)(bank.size() - 2) / (remaining + 1); - for (int i = 0; i < remaining; ++i) { - int idx = 1 + (int)((i + 1) * step); - idx = std::min(idx, (int)bank.size() - 2); - selected.push_back(idx); - } - } - std::sort(selected.begin(), selected.end()); - selected.erase(std::unique(selected.begin(), selected.end()), selected.end()); - return selected; -} - -// Compute mask IoU between two binary masks. -static float sam3_mask_iou(const uint8_t* a, const uint8_t* b, int n) { - int inter = 0, uni = 0; - for (int i = 0; i < n; ++i) { - bool va = a[i] > 127; - bool vb = b[i] > 127; - if (va && vb) ++inter; - if (va || vb) ++uni; - } - return (uni > 0) ? (float)inter / uni : 0.0f; -} - -/***************************************************************************** -** Post-processing: hole filling and sprinkle removal (Phase 7, Step 7.8) -*****************************************************************************/ - -// Fill small holes in a binary mask using BFS connected components. -static void sam3_fill_holes(uint8_t* mask, int w, int h, int area_threshold) { - const int n = w * h; - std::vector labels(n, -1); - int next_label = 0; - std::vector component_sizes; - - for (int i = 0; i < n; ++i) { - if (mask[i] > 127 || labels[i] >= 0) continue; - int label = next_label++; - component_sizes.push_back(0); - std::vector queue; - queue.push_back(i); - labels[i] = label; - int head = 0; - bool touches_border = false; - while (head < (int)queue.size()) { - int p = queue[head++]; - component_sizes[label]++; - int px = p % w, py = p / w; - if (px == 0 || px == w - 1 || py == 0 || py == h - 1) touches_border = true; - int dx[] = {-1, 1, 0, 0}; - int dy[] = {0, 0, -1, 1}; - for (int d = 0; d < 4; ++d) { - int nx2 = px + dx[d], ny2 = py + dy[d]; - if (nx2 < 0 || nx2 >= w || ny2 < 0 || ny2 >= h) continue; - int ni = ny2 * w + nx2; - if (mask[ni] <= 127 && labels[ni] < 0) { - labels[ni] = label; - queue.push_back(ni); - } - } - } - if (touches_border) component_sizes[label] = area_threshold + 1; - } - for (int i = 0; i < n; ++i) { - if (mask[i] <= 127 && labels[i] >= 0 && component_sizes[labels[i]] <= area_threshold) { - mask[i] = 255; - } - } -} - -// Remove small foreground sprinkles. -static void sam3_remove_sprinkles(uint8_t* mask, int w, int h, int area_threshold) { - const int n = w * h; - std::vector labels(n, -1); - int next_label = 0; - std::vector component_sizes; - - for (int i = 0; i < n; ++i) { - if (mask[i] <= 127 || labels[i] >= 0) continue; - int label = next_label++; - component_sizes.push_back(0); - std::vector queue; - queue.push_back(i); - labels[i] = label; - int head = 0; - while (head < (int)queue.size()) { - int p = queue[head++]; - component_sizes[label]++; - int px = p % w, py = p / w; - int dx[] = {-1, 1, 0, 0}; - int dy[] = {0, 0, -1, 1}; - for (int d = 0; d < 4; ++d) { - int nx2 = px + dx[d], ny2 = py + dy[d]; - if (nx2 < 0 || nx2 >= w || ny2 < 0 || ny2 >= h) continue; - int ni = ny2 * w + nx2; - if (mask[ni] > 127 && labels[ni] < 0) { - labels[ni] = label; - queue.push_back(ni); - } - } - } - } - for (int i = 0; i < n; ++i) { - if (mask[i] > 127 && labels[i] >= 0 && component_sizes[labels[i]] <= area_threshold) { - mask[i] = 0; - } - } -} - -// Resolve overlapping masks: higher-scoring instances take priority. -static void sam3_resolve_overlaps(std::vector& dets) { - if (dets.size() <= 1) return; - const int w = dets[0].mask.width; - const int h = dets[0].mask.height; - if (w == 0 || h == 0) return; - std::sort(dets.begin(), dets.end(), [](const sam3_detection& a, const sam3_detection& b) { - return a.score > b.score; - }); - const int n = w * h; - for (int i = 0; i < n; ++i) { - bool claimed = false; - for (auto& det : dets) { - if (det.mask.data.empty()) continue; - if (claimed) { - det.mask.data[i] = 0; - } else if (det.mask.data[i] > 127) { - claimed = true; - } - } - } -} - -/***************************************************************************** -** Post-processing: NMS, bilinear interpolation, mask binarization -*****************************************************************************/ - -// Compute IoU between two boxes [x0, y0, x1, y1]. -static float sam3_box_iou(const sam3_box& a, const sam3_box& b) { - float x0 = std::max(a.x0, b.x0); - float y0 = std::max(a.y0, b.y0); - float x1 = std::min(a.x1, b.x1); - float y1 = std::min(a.y1, b.y1); - float inter = std::max(0.0f, x1 - x0) * std::max(0.0f, y1 - y0); - float area_a = (a.x1 - a.x0) * (a.y1 - a.y0); - float area_b = (b.x1 - b.x0) * (b.y1 - b.y0); - float uni = area_a + area_b - inter; - return (uni > 0.0f) ? inter / uni : 0.0f; -} - -// Non-maximum suppression on detections, sorted by score descending. -// Returns indices of kept detections. -static std::vector sam3_nms(const std::vector& dets, float iou_thresh) { - // Sort indices by score descending - std::vector indices(dets.size()); - for (int i = 0; i < (int)dets.size(); ++i) indices[i] = i; - std::sort(indices.begin(), indices.end(), [&](int a, int b) { - return dets[a].score > dets[b].score; - }); - - std::vector suppressed(dets.size(), false); - std::vector keep; - - for (int idx : indices) { - if (suppressed[idx]) continue; - keep.push_back(idx); - for (int j : indices) { - if (suppressed[j] || j == idx) continue; - if (sam3_box_iou(dets[idx].box, dets[j].box) > iou_thresh) { - suppressed[j] = true; - } - } - } - - return keep; -} - -// Bilinear interpolation of a flat mask [H_in * W_in] to [H_out * W_out]. -// Uses double for coordinate math to match PyTorch F.interpolate precision. -static std::vector sam3_bilinear_interpolate(const float* src, int src_w, int src_h, - int dst_w, int dst_h) { - std::vector dst(dst_w * dst_h); - const double sx = (double)src_w / dst_w; - const double sy = (double)src_h / dst_h; - - for (int y = 0; y < dst_h; ++y) { - double fy = (y + 0.5) * sy - 0.5; - fy = std::max(0.0, std::min(fy, (double)(src_h - 1))); - const int y0 = std::min((int)fy, src_h - 2); - const int y1 = y0 + 1; - const float wy = (float)(fy - y0); - - for (int x = 0; x < dst_w; ++x) { - double fx = (x + 0.5) * sx - 0.5; - fx = std::max(0.0, std::min(fx, (double)(src_w - 1))); - const int x0 = std::min((int)fx, src_w - 2); - const int x1 = x0 + 1; - const float wx = (float)(fx - x0); - - float v = (1 - wy) * ((1 - wx) * src[y0 * src_w + x0] + wx * src[y0 * src_w + x1]) + wy * ((1 - wx) * src[y1 * src_w + x0] + wx * src[y1 * src_w + x1]); - dst[y * dst_w + x] = v; - } - } - return dst; -} - -// Convert (cx, cy, w, h) in [0,1] to (x0, y0, x1, y1) in pixel coordinates. -static sam3_box sam3_cxcywh_to_xyxy(float cx, float cy, float w, float h, - int img_w, int img_h) { - sam3_box box; - box.x0 = (cx - w * 0.5f) * img_w; - box.y0 = (cy - h * 0.5f) * img_h; - box.x1 = (cx + w * 0.5f) * img_w; - box.y1 = (cy + h * 0.5f) * img_h; - // Clamp to image bounds - box.x0 = std::max(0.0f, std::min(box.x0, (float)img_w)); - box.y0 = std::max(0.0f, std::min(box.y0, (float)img_h)); - box.x1 = std::max(0.0f, std::min(box.x1, (float)img_w)); - box.y1 = std::max(0.0f, std::min(box.y1, (float)img_h)); - return box; -} - -/***************************************************************************** -** Image segmentation — PCS (text-prompted) -*****************************************************************************/ - -sam3_result sam3_segment_pcs(sam3_state& state, - const sam3_model& model, - const sam3_pcs_params& params) { - if (model.hparams.visual_only || model.hparams.is_sam2()) { - fprintf(stderr, "%s: ERROR: PCS not available on %s model\n", - __func__, model.hparams.is_sam2() ? "SAM2" : "visual-only"); - return sam3_result{}; - } - -#if SAM3_LOG_LEVEL >= 1 - auto t_start = std::chrono::high_resolution_clock::now(); -#endif - const auto& hp = model.hparams; - const int D = hp.neck_dim; // 256 - const int H = hp.n_img_embd(); // 72 - const int L = hp.text_ctx_len; // 32 - const int NQ = hp.ddec_num_queries; // 200 - const int N_spatial = H * H; // 5184 - sam3_result result; - - // ── Check that image has been encoded ──────────────────────────────── - if (!state.neck_det[0]) { - fprintf(stderr, "%s: image not encoded — call sam3_encode_image first\n", __func__); - return result; - } - - // ── Tokenize text prompt ───────────────────────────────────────────── - auto token_ids = sam3_tokenize(const_cast(model.tokenizer), - params.text_prompt, L); - if (token_ids.empty()) { - fprintf(stderr, "%s: failed to tokenize text prompt\n", __func__); - return result; - } - - SAM3_LOG(2, "%s: text='%s', %zu tokens\n", __func__, - params.text_prompt.c_str(), token_ids.size()); - - // ── Helper: run a sub-graph with its own context and allocator ────── - // Each stage below follows this exact pattern: - // 1. Create fresh ggml_context + graph + allocator - // 2. Create INPUT tensors (ggml_set_input) - // 3. Build the computation graph - // 4. Mark outputs (ggml_set_output) - // 5. Allocate → set input data → compute → read output data - // 6. Free allocator + context - // This ensures ZERO buffer sharing between stages. - - // ── Pre-compute shared CPU data used by multiple stages ───────────── - const int n_boxes = (int)(params.pos_exemplars.size() + params.neg_exemplars.size()); - const int N_geo = n_boxes + 1; // +1 for CLS - const int T = L + N_geo; // total prompt tokens (text + geometry) - - // Read image features and PE from state into CPU buffers (used by multiple stages) - std::vector img_feats_cpu(D * N_spatial); - std::vector img_pe_cpu(D * N_spatial); - ggml_backend_tensor_get(state.neck_det[2], img_feats_cpu.data(), 0, - D * N_spatial * sizeof(float)); - ggml_backend_tensor_get(state.neck_det_pe[2], img_pe_cpu.data(), 0, - D * N_spatial * sizeof(float)); - - // Pre-compute constant input vectors - std::vector sine_dim_t_cpu(64); - for (int i = 0; i < 64; ++i) - sine_dim_t_cpu[i] = 2.0f * (float)M_PI / powf(10000.0f, 2.0f * (float)i / 128.0f); - - std::vector rpb_coords_cpu(H); - for (int i = 0; i < H; ++i) rpb_coords_cpu[i] = (float)i / (float)H; - - // Build combined attention bias for the prompt: [T] = [L + N_geo] - // 0.0 for valid tokens, -1e9 for padding text tokens, 0.0 for geo tokens - std::vector combined_bias_cpu(T); - for (int i = 0; i < L; ++i) - combined_bias_cpu[i] = (token_ids[i] != 0) ? 0.0f : -1.0e9f; - for (int i = 0; i < N_geo; ++i) - combined_bias_cpu[L + i] = 0.0f; - - // Build text validity mask for DotProductScoring: [T] - int n_valid_tokens = 0; - for (int i = 0; i < L; ++i) - if (token_ids[i] != 0) ++n_valid_tokens; - n_valid_tokens += N_geo; - if (n_valid_tokens == 0) n_valid_tokens = 1; - float tvm_scale = (float)T / (float)n_valid_tokens; - std::vector text_valid_mask_cpu(T); - for (int i = 0; i < L; ++i) - text_valid_mask_cpu[i] = (token_ids[i] != 0) ? tvm_scale : 0.0f; - for (int i = 0; i < N_geo; ++i) - text_valid_mask_cpu[L + i] = tvm_scale; - - /* - ** ── SUB-GRAPH 1: Text Encoder ──────────────────────────────────── - */ - std::vector text_feats_cpu(D * L); - { - const size_t sz = ggml_tensor_overhead() * 16384 + ggml_graph_overhead() * 2; - struct ggml_init_params gp = {sz, nullptr, true}; - auto* ctx = ggml_init(gp); - - auto* inp = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, L); - ggml_set_name(inp, "text_token_ids"); - ggml_set_input(inp); - - auto* out = sam3_build_text_encoder_graph(ctx, inp, model); - ggml_set_output(out); - - auto* graph = ggml_new_graph_custom(ctx, 16384, false); - ggml_build_forward_expand(graph, out); - - auto* causal = ggml_get_tensor(ctx, "causal_mask"); - auto* alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); - if (!ggml_gallocr_reserve(alloc, graph) || !ggml_gallocr_alloc_graph(alloc, graph)) { - fprintf(stderr, "%s: text encoder alloc failed\n", __func__); - ggml_gallocr_free(alloc); - ggml_free(ctx); - return result; - } - - if (inp->buffer) { - ggml_backend_tensor_set(inp, token_ids.data(), 0, L * sizeof(int32_t)); - } else { - fprintf(stderr, "%s: ERROR: text encoder input tensor has no buffer!\n", __func__); - ggml_gallocr_free(alloc); - ggml_free(ctx); - return result; - } - if (causal && causal->buffer) { - std::vector cm(L * L); - sam3_fill_causal_mask(cm.data(), L); - ggml_backend_tensor_set(causal, cm.data(), 0, L * L * sizeof(ggml_fp16_t)); - } - - if (!sam3_graph_compute(model.backend, graph, state.n_threads)) { - ggml_gallocr_free(alloc); - ggml_free(ctx); - return result; - } - ggml_backend_tensor_get(out, text_feats_cpu.data(), 0, D * L * sizeof(float)); - - ggml_gallocr_free(alloc); - ggml_free(ctx); - } - - SAM3_LOG(2, "%s: text encoder done\n", __func__); - - /* - ** ── SUB-GRAPH 2: Geometry Encoder ──────────────────────────────── - */ - std::vector geo_feats_cpu(D * N_geo); - { - const size_t sz = ggml_tensor_overhead() * 4096 + ggml_graph_overhead() * 2; - struct ggml_init_params gp = {sz, nullptr, true}; - auto* ctx = ggml_init(gp); - - auto* g_img = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, N_spatial, 1); - ggml_set_name(g_img, "geo_img"); - ggml_set_input(g_img); - auto* g_pe = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, N_spatial, 1); - ggml_set_name(g_pe, "geo_pe"); - ggml_set_input(g_pe); - - auto gr = sam3_build_geom_enc_graph(ctx, model, params, g_img, g_pe); - ggml_set_output(gr.geo_feats); - - auto* graph = ggml_new_graph_custom(ctx, 4096, false); - ggml_build_forward_expand(graph, gr.geo_feats); - - auto* alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); - if (!ggml_gallocr_reserve(alloc, graph) || !ggml_gallocr_alloc_graph(alloc, graph)) { - fprintf(stderr, "%s: geometry encoder alloc failed\n", __func__); - ggml_gallocr_free(alloc); - ggml_free(ctx); - return result; - } - - ggml_backend_tensor_set(g_img, img_feats_cpu.data(), 0, D * N_spatial * sizeof(float)); - ggml_backend_tensor_set(g_pe, img_pe_cpu.data(), 0, D * N_spatial * sizeof(float)); - - // Pre-transformer input (CLS + optional box embeddings) - { - const float* feats_ptr = n_boxes > 0 ? img_feats_cpu.data() : nullptr; - auto geom_data = sam3_precompute_geom_input(model, params, feats_ptr, H, H); - auto* gi = ggml_get_tensor(ctx, "geom_post_final_proj"); - if (gi) ggml_backend_tensor_set(gi, geom_data.data(), 0, geom_data.size() * sizeof(float)); - } - - if (!sam3_graph_compute(model.backend, graph, state.n_threads)) { - ggml_gallocr_free(alloc); - ggml_free(ctx); - return result; - } - ggml_backend_tensor_get(gr.geo_feats, geo_feats_cpu.data(), 0, - D * N_geo * sizeof(float)); - - ggml_gallocr_free(alloc); - ggml_free(ctx); - } - - SAM3_LOG(2, "%s: geometry encoder done\n", __func__); - - // Build combined prompt on CPU: [text_feats, geo_feats] → [D * T] - std::vector combined_prompt_cpu(D * T); - memcpy(combined_prompt_cpu.data(), text_feats_cpu.data(), D * L * sizeof(float)); - memcpy(combined_prompt_cpu.data() + D * L, geo_feats_cpu.data(), D * N_geo * sizeof(float)); - - SAM3_LOG(2, "%s: starting fusion encoder\n", __func__); - /* - ** ── SUB-GRAPH 3: Fusion Encoder ────────────────────────────────── - */ - std::vector fenc_output_cpu(D * N_spatial); - { - const size_t sz = ggml_tensor_overhead() * 16384 + ggml_graph_overhead() * 2; - struct ggml_init_params gp = {sz, nullptr, true}; - auto* ctx = ggml_init(gp); - - auto* img = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, N_spatial, 1); - ggml_set_name(img, "fenc_img"); - ggml_set_input(img); - auto* pe = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, N_spatial, 1); - ggml_set_name(pe, "fenc_pe"); - ggml_set_input(pe); - auto* prompt = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, T, 1); - ggml_set_name(prompt, "fenc_prompt"); - ggml_set_input(prompt); - auto* bias = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, T, 1, 1); - ggml_set_name(bias, "fenc_bias"); - ggml_set_input(bias); - - auto* out = sam3_build_fenc_graph(ctx, model, img, prompt, pe, bias); - ggml_set_output(out); - - auto* graph = ggml_new_graph_custom(ctx, 16384, false); - ggml_build_forward_expand(graph, out); - - auto* alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); - if (!ggml_gallocr_reserve(alloc, graph) || !ggml_gallocr_alloc_graph(alloc, graph)) { - fprintf(stderr, "%s: fusion encoder alloc failed\n", __func__); - ggml_gallocr_free(alloc); - ggml_free(ctx); - return result; - } - - ggml_backend_tensor_set(img, img_feats_cpu.data(), 0, D * N_spatial * sizeof(float)); - ggml_backend_tensor_set(pe, img_pe_cpu.data(), 0, D * N_spatial * sizeof(float)); - ggml_backend_tensor_set(prompt, combined_prompt_cpu.data(), 0, D * T * sizeof(float)); - ggml_backend_tensor_set(bias, combined_bias_cpu.data(), 0, T * sizeof(float)); - - if (!sam3_graph_compute(model.backend, graph, state.n_threads)) { - ggml_gallocr_free(alloc); - ggml_free(ctx); - return result; - } - ggml_backend_tensor_get(out, fenc_output_cpu.data(), 0, D * N_spatial * sizeof(float)); - - SAM3_LOG(2, "%s: fenc_out[0..4] = [%.6f, %.6f, %.6f, %.6f, %.6f]\n", - __func__, fenc_output_cpu[0], fenc_output_cpu[1], fenc_output_cpu[2], - fenc_output_cpu[3], fenc_output_cpu[4]); - - ggml_gallocr_free(alloc); - ggml_free(ctx); - } - - SAM3_LOG(2, "%s: fusion encoder done\n", __func__); - /* - ** ── SUB-GRAPH 4: DETR Decoder + Scoring ────────────────────────── - */ - std::vector scores_data(NQ); - std::vector boxes_data(4 * NQ); - std::vector queries_data(D * (NQ + 1)); // 201 = NQ + presence token - float presence_logit = 0.0f; - { - const size_t sz = ggml_tensor_overhead() * 32768 + ggml_graph_overhead() * 2; - struct ggml_init_params gp = {sz, nullptr, true}; - auto* ctx = ggml_init(gp); - - auto* enc = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, N_spatial, 1); - ggml_set_name(enc, "ddec_enc"); - ggml_set_input(enc); - auto* pe = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, N_spatial, 1); - ggml_set_name(pe, "ddec_pe"); - ggml_set_input(pe); - auto* txt = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, T, 1); - ggml_set_name(txt, "ddec_text"); - ggml_set_input(txt); - auto* sdt = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, 64); - ggml_set_name(sdt, "sine_dim_t"); - ggml_set_input(sdt); - auto* rpb = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, H); - ggml_set_name(rpb, "rpb_coords"); - ggml_set_input(rpb); - auto* tab = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, T, 1, 1); - ggml_set_name(tab, "text_attn_bias"); - ggml_set_input(tab); - auto* tvm = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, T, 1, 1); - ggml_set_name(tvm, "text_valid_mask"); - ggml_set_input(tvm); - - auto dout = sam3_build_ddec_graph(ctx, model, enc, pe, txt, sdt, rpb, tab, tvm); - ggml_set_output(dout.class_scores); - ggml_set_output(dout.pred_boxes); - ggml_set_output(dout.presence_score); - ggml_set_output(dout.queries); - - auto* graph = ggml_new_graph_custom(ctx, 65536, false); - ggml_build_forward_expand(graph, dout.class_scores); - ggml_build_forward_expand(graph, dout.pred_boxes); - ggml_build_forward_expand(graph, dout.presence_score); - ggml_build_forward_expand(graph, dout.queries); - - auto* alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); - if (!ggml_gallocr_reserve(alloc, graph) || !ggml_gallocr_alloc_graph(alloc, graph)) { - fprintf(stderr, "%s: DETR decoder alloc failed\n", __func__); - ggml_gallocr_free(alloc); - ggml_free(ctx); - return result; - } - - ggml_backend_tensor_set(enc, fenc_output_cpu.data(), 0, D * N_spatial * sizeof(float)); - ggml_backend_tensor_set(pe, img_pe_cpu.data(), 0, D * N_spatial * sizeof(float)); - ggml_backend_tensor_set(txt, combined_prompt_cpu.data(), 0, D * T * sizeof(float)); - ggml_backend_tensor_set(sdt, sine_dim_t_cpu.data(), 0, 64 * sizeof(float)); - ggml_backend_tensor_set(rpb, rpb_coords_cpu.data(), 0, H * sizeof(float)); - ggml_backend_tensor_set(tab, combined_bias_cpu.data(), 0, T * sizeof(float)); - ggml_backend_tensor_set(tvm, text_valid_mask_cpu.data(), 0, T * sizeof(float)); - - // Presence token position encoding: zeros - auto* qpos = ggml_get_tensor(ctx, "ddec_query_pos_pres"); - if (qpos) { - std::vector z(D, 0.0f); - ggml_backend_tensor_set(qpos, z.data(), 0, D * sizeof(float)); - } - auto* rpbz = ggml_get_tensor(ctx, "rpb_pres_zeros"); - if (rpbz) { - int n = (int)(rpbz->ne[0] * rpbz->ne[1] * rpbz->ne[2] * rpbz->ne[3]); - std::vector z(n, 0.0f); - ggml_backend_tensor_set(rpbz, z.data(), 0, n * sizeof(float)); - } - - if (!sam3_graph_compute(model.backend, graph, state.n_threads)) { - ggml_gallocr_free(alloc); - ggml_free(ctx); - return result; - } - - ggml_backend_tensor_get(dout.class_scores, scores_data.data(), 0, NQ * sizeof(float)); - ggml_backend_tensor_get(dout.pred_boxes, boxes_data.data(), 0, 4 * NQ * sizeof(float)); - ggml_backend_tensor_get(dout.presence_score, &presence_logit, 0, sizeof(float)); - ggml_backend_tensor_get(dout.queries, queries_data.data(), 0, - D * (NQ + 1) * sizeof(float)); - - ggml_gallocr_free(alloc); - ggml_free(ctx); - } - - float presence_prob = 1.0f / (1.0f + expf(-presence_logit)); - - SAM3_LOG(2, "%s: DETR decoder done\n", __func__); - /* - ** ── SUB-GRAPH 5: Segmentation Head ─────────────────────────────── - */ - const int mask_hw = H * 4; // 288 for SAM3 - std::vector all_masks(NQ * mask_hw * mask_hw); - { - const size_t sz = ggml_tensor_overhead() * 16384 + ggml_graph_overhead() * 2; - struct ggml_init_params gp = {sz, nullptr, true}; - auto* ctx = ggml_init(gp); - - auto* enc_h = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, N_spatial, 1); - ggml_set_name(enc_h, "seg_enc"); - ggml_set_input(enc_h); - - // FPN features at 3 scales — create fresh inputs - const int H0 = H * 4, H1 = H * 2; - auto* fpn0 = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, D, H0, H0, 1); - ggml_set_name(fpn0, "seg_fpn0"); - ggml_set_input(fpn0); - auto* fpn1 = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, D, H1, H1, 1); - ggml_set_name(fpn1, "seg_fpn1"); - ggml_set_input(fpn1); - auto* fpn2 = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, D, H, H, 1); - ggml_set_name(fpn2, "seg_fpn2"); - ggml_set_input(fpn2); - struct ggml_tensor* fpn_feats[3] = {fpn0, fpn1, fpn2}; - - // Object queries (skip presence token at index 0 → start at index 1) - auto* oq = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, NQ, 1); - ggml_set_name(oq, "seg_queries"); - ggml_set_input(oq); - - auto* txt = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, T, 1); - ggml_set_name(txt, "seg_text"); - ggml_set_input(txt); - auto* tab = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, T, 1, 1); - ggml_set_name(tab, "seg_bias"); - ggml_set_input(tab); - - auto* out = sam3_build_seg_head_graph(ctx, model, enc_h, fpn_feats, - oq, txt, tab); - ggml_set_output(out); - - auto* graph = ggml_new_graph_custom(ctx, 32768, false); - ggml_build_forward_expand(graph, out); - - SAM3_LOG(2, "%s: seg head graph: %d nodes\n", __func__, ggml_graph_n_nodes(graph)); - - auto* alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); - if (!ggml_gallocr_reserve(alloc, graph) || !ggml_gallocr_alloc_graph(alloc, graph)) { - fprintf(stderr, "%s: segmentation head alloc failed\n", __func__); - ggml_gallocr_free(alloc); - ggml_free(ctx); - return result; - } - - ggml_backend_tensor_set(enc_h, fenc_output_cpu.data(), 0, - D * N_spatial * sizeof(float)); - - // Copy FPN features from state - { - size_t s0 = (size_t)D * H0 * H0 * sizeof(float); - size_t s1 = (size_t)D * H1 * H1 * sizeof(float); - size_t s2 = (size_t)D * H * H * sizeof(float); - std::vector b0(D * H0 * H0), b1(D * H1 * H1), b2(D * H * H); - ggml_backend_tensor_get(state.neck_det[0], b0.data(), 0, s0); - if (fpn0->buffer) ggml_backend_tensor_set(fpn0, b0.data(), 0, s0); - ggml_backend_tensor_get(state.neck_det[1], b1.data(), 0, s1); - if (fpn1->buffer) ggml_backend_tensor_set(fpn1, b1.data(), 0, s1); - ggml_backend_tensor_get(state.neck_det[2], b2.data(), 0, s2); - if (fpn2->buffer) ggml_backend_tensor_set(fpn2, b2.data(), 0, s2); - } - - // Object queries: extract from DETR queries (skip presence token at slot 0) - // queries_data is flat [D * 201], presence token is at positions [0..D-1] - // Object queries start at position [D..D*(NQ+1)-1] - ggml_backend_tensor_set(oq, queries_data.data() + D, 0, D * NQ * sizeof(float)); - - ggml_backend_tensor_set(txt, combined_prompt_cpu.data(), 0, D * T * sizeof(float)); - ggml_backend_tensor_set(tab, combined_bias_cpu.data(), 0, T * sizeof(float)); - - if (!sam3_graph_compute(model.backend, graph, state.n_threads)) { - ggml_gallocr_free(alloc); - ggml_free(ctx); - return result; - } - ggml_backend_tensor_get(out, all_masks.data(), 0, all_masks.size() * sizeof(float)); - - ggml_gallocr_free(alloc); - ggml_free(ctx); - } - - /* - ** ── Post-processing: thresholding + NMS + mask resize ──────────── - */ - std::vector dets; - for (int q = 0; q < NQ; ++q) { - float class_prob = 1.0f / (1.0f + expf(-scores_data[q])); - float score = class_prob * presence_prob; - if (score < params.score_threshold) continue; - - sam3_detection det; - float cx = boxes_data[0 + q * 4]; - float cy = boxes_data[1 + q * 4]; - float bw = boxes_data[2 + q * 4]; - float bh = boxes_data[3 + q * 4]; - - det.box = sam3_cxcywh_to_xyxy(cx, cy, bw, bh, state.orig_width, state.orig_height); - det.score = score; - - const float* mask_ptr = all_masks.data() + q * mask_hw * mask_hw; - auto mask_resized = sam3_bilinear_interpolate(mask_ptr, mask_hw, mask_hw, - state.orig_width, state.orig_height); - det.mask.width = state.orig_width; - det.mask.height = state.orig_height; - det.mask.data.resize(state.orig_width * state.orig_height); - for (int i = 0; i < (int)mask_resized.size(); ++i) - det.mask.data[i] = (mask_resized[i] > 0.0f) ? 255 : 0; - det.mask.iou_score = score; - - dets.push_back(std::move(det)); - } - - SAM3_LOG(2, "%s: %zu detections above threshold %.2f (presence=%.3f, logit=%.3f)\n", - __func__, dets.size(), params.score_threshold, presence_prob, presence_logit); - - auto keep = sam3_nms(dets, params.nms_threshold); - for (int i = 0; i < (int)keep.size(); ++i) { - dets[keep[i]].instance_id = i + 1; - result.detections.push_back(std::move(dets[keep[i]])); - } - - SAM3_LOG(2, "%s: %zu detections after NMS\n", __func__, result.detections.size()); - -#if SAM3_LOG_LEVEL >= 1 - auto t_end = std::chrono::high_resolution_clock::now(); - double total_ms = std::chrono::duration(t_end - t_start).count(); - SAM3_LOG(1, "%s: completed in %.1f ms\n", __func__, total_ms); -#endif - - return result; -} - -/***************************************************************************** -** SAM attention helper (separate Q, K, V weight/bias) -*****************************************************************************/ - -static struct ggml_tensor* sam3_sam_attention( - struct ggml_context* ctx, - struct ggml_tensor* q_in, // [D, N_q, B] - struct ggml_tensor* k_in, // [D, N_kv, B] - struct ggml_tensor* v_in, // [D, N_kv, B] - const sam3_sam_attn& attn, - int n_heads) { - const int64_t N_q = q_in->ne[1]; - const int64_t B = q_in->ne[2]; - const int64_t N_kv = k_in->ne[1]; - - // Project - auto* Q = ggml_add(ctx, ggml_mul_mat(ctx, attn.q_w, q_in), attn.q_b); - auto* K = ggml_add(ctx, ggml_mul_mat(ctx, attn.k_w, k_in), attn.k_b); - auto* V = ggml_add(ctx, ggml_mul_mat(ctx, attn.v_w, v_in), attn.v_b); - - // Debug: mark projections for the first SA call (N_q=8 tokens, block 0) - static int _sa_call_count = 0; - if (_sa_call_count == 0 && N_q <= 16) { - ggml_set_name(Q, "dbg_sa0_Q_proj"); - ggml_set_output(Q); - ggml_set_name(V, "dbg_sa0_V_proj"); - ggml_set_output(V); - } - _sa_call_count++; - - // internal_dim = out_proj cols = attn.q_w->ne[1] - const int64_t ID = attn.q_w->ne[1]; - const int64_t HD = ID / n_heads; - - // Reshape to multi-head: [HD, N, NH, B] - Q = ggml_reshape_4d(ctx, Q, HD, n_heads, N_q, B); - Q = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 1, 3)); // [HD, N_q, NH, B] - - K = ggml_reshape_4d(ctx, K, HD, n_heads, N_kv, B); - K = ggml_cont(ctx, ggml_permute(ctx, K, 0, 2, 1, 3)); // [HD, N_kv, NH, B] - - V = ggml_reshape_4d(ctx, V, HD, n_heads, N_kv, B); - V = ggml_cont(ctx, ggml_permute(ctx, V, 0, 2, 1, 3)); // [HD, N_kv, NH, B] contiguous - - // Attention - float scale = 1.0f / sqrtf((float)HD); - auto* out = ggml_flash_attn_ext(ctx, Q, K, V, nullptr, scale, 0.0f, 0.0f); - // out: [HD, NH, N_q, B] (flash_attn_ext swaps dims 1,2 vs input) - -#if 0 // Manual SDPA (for debugging only) - auto* Q3 = ggml_reshape_3d(ctx, Q, HD, N_q, n_heads * B); - auto* K3 = ggml_reshape_3d(ctx, K, HD, N_kv, n_heads * B); - auto* V3 = ggml_reshape_3d(ctx, V, HD, N_kv, n_heads * B); - // QK^T: ggml_mul_mat(K, Q) → K^T @ Q → [N_kv, N_q, NH*B] - auto* attn_scores = ggml_mul_mat(ctx, K3, Q3); - attn_scores = ggml_scale(ctx, attn_scores, scale); - attn_scores = ggml_soft_max(ctx, attn_scores); - - // attn @ V: need attn^T [N_q, N_kv] and V^T [HD, N_kv] - // ggml_mul_mat(attn^T, V) = (attn^T)^T @ V = attn @ V = [N_q, HD]... no. - // ggml_mul_mat(A, B) = A^T @ B where A=[K, M], B=[K, N] → [M, N] - // Want: output[q, d] = sum_k attn[q, k] * V[k, d] - // = (V^T @ attn^T)^T... let me think differently. - // attn_scores is [N_kv, N_q, NH*B]. For each head: - // attn[k, q] = attn_scores[k, q] (col q has the weights for query q) - // V3 is [HD, N_kv, NH*B]. - // Want: out[d, q] = sum_k V[d, k] * attn[k, q] = V @ attn - // = ggml_mul_mat? mul_mat(A, B) = A^T B with A=[K, M], B=[K, N] → [M, N] - // V has ne=[HD, N_kv, ...]. attn has ne=[N_kv, N_q, ...]. - // If A=V3 (ne0=HD, ne1=N_kv) and B=attn_scores (ne0=N_kv, ne1=N_q): - // Shared dim ne0: V3 ne0=HD ≠ attn ne0=N_kv. Mismatch! - // - // Need to transpose V: V^T is [N_kv, HD]. Then A=V^T, B=attn_scores. - // A ne0=N_kv, B ne0=N_kv → shared. A^T B = V @ attn → [HD, N_q]. ✓ - auto* VT = ggml_permute(ctx, V3, 1, 0, 2, 3); // [N_kv, HD, NH*B] - VT = ggml_cont(ctx, VT); - auto* out3 = ggml_mul_mat(ctx, VT, attn_scores); // [HD, N_q, NH*B] - - // Reshape back to 4D: [HD, N_q, NH, B] - auto* out = ggml_reshape_4d(ctx, out3, HD, N_q, n_heads, B); - // Permute to [HD, NH, N_q, B] to match flash_attn_ext output convention - out = ggml_cont(ctx, ggml_permute(ctx, out, 0, 2, 1, 3)); -#endif - - // Merge heads: [ID=HD*NH, N_q, B] - auto* merged = ggml_reshape_3d(ctx, out, ID, N_q, B); - - // Debug: mark merged attention output for first SA call - static int _sa_merge_count = 0; - if (_sa_merge_count == 0 && N_q <= 16) { - ggml_set_name(merged, "dbg_sa0_merged"); - ggml_set_output(merged); - } - _sa_merge_count++; - - // Output projection - out = ggml_mul_mat(ctx, attn.out_w, merged); - out = ggml_add(ctx, out, attn.out_b); - - return out; -} - -/***************************************************************************** -** SAM prompt encoder — graph building (Phase 6, Step 6.1) -*****************************************************************************/ - -// Random Fourier positional encoding for a single (x, y) coordinate -// coords_norm: normalized to [0, 1], pe_gaussian: PyTorch row-major [2, 128] -// Output: [256] = [sin(128); cos(128)] -static void sam3_pe_encode_coord(float* out, float x_norm, float y_norm, - const float* pe_gauss, int num_pos_feats) { - // Map [0,1] → [-1,1] - float coords[2] = {2.0f * x_norm - 1.0f, 2.0f * y_norm - 1.0f}; - - // coords @ pe_gaussian → [128] - for (int i = 0; i < num_pos_feats; ++i) { - float dot = coords[0] * pe_gauss[i] + - coords[1] * pe_gauss[num_pos_feats + i]; - dot *= 2.0f * (float)M_PI; - out[i] = sinf(dot); - out[i + num_pos_feats] = cosf(dot); - } -} - -// Read SAM prompt encoder weights from GPU and cache them in state. -// Also pre-computes the dense PE grid and no-mask tiled embedding. -// These never change between PVS calls for the same model. -static void sam3_populate_pe_cache(sam3_state& state, const sam3_model& model) { - if (state.pe_cache_valid) return; - - const int D = model.hparams.sam_embed_dim; // 256 - const int H = sam3_eff_feat_size(state, model.hparams); - const int num_pos_feats = D / 2; // 128 - const int pe_nel = 2 * num_pos_feats; // 256 - const auto& pe = model.sam_pe; - - state.pe_gauss_cache.resize(pe_nel); - if (pe.pe_gaussian->type == GGML_TYPE_F16) { - std::vector tmp(pe_nel); - ggml_backend_tensor_get(pe.pe_gaussian, tmp.data(), 0, pe_nel * sizeof(ggml_fp16_t)); - ggml_fp16_to_fp32_row(tmp.data(), state.pe_gauss_cache.data(), pe_nel); - } else { - ggml_backend_tensor_get(pe.pe_gaussian, state.pe_gauss_cache.data(), 0, pe_nel * sizeof(float)); - } - - for (int i = 0; i < 4; ++i) { - if (pe.point_embed[i]->type == GGML_TYPE_F16) { - std::vector tmp(D); - ggml_backend_tensor_get(pe.point_embed[i], tmp.data(), 0, D * sizeof(ggml_fp16_t)); - ggml_fp16_to_fp32_row(tmp.data(), state.point_emb_cache[i], D); - } else { - ggml_backend_tensor_get(pe.point_embed[i], state.point_emb_cache[i], 0, D * sizeof(float)); - } - } - - if (pe.not_a_point_embed->type == GGML_TYPE_F16) { - std::vector tmp(D); - ggml_backend_tensor_get(pe.not_a_point_embed, tmp.data(), 0, D * sizeof(ggml_fp16_t)); - ggml_fp16_to_fp32_row(tmp.data(), state.not_a_point_cache, D); - } else { - ggml_backend_tensor_get(pe.not_a_point_embed, state.not_a_point_cache, 0, D * sizeof(float)); - } - - if (pe.no_mask_embed->type == GGML_TYPE_F16) { - std::vector tmp(D); - ggml_backend_tensor_get(pe.no_mask_embed, tmp.data(), 0, D * sizeof(ggml_fp16_t)); - ggml_fp16_to_fp32_row(tmp.data(), state.no_mask_emb_cache, D); - } else { - ggml_backend_tensor_get(pe.no_mask_embed, state.no_mask_emb_cache, 0, D * sizeof(float)); - } - - state.dense_pe_cache.resize(D * H * H); - for (int row = 0; row < H; ++row) { - for (int col = 0; col < H; ++col) { - float x_norm = ((float)col + 0.5f) / (float)H; - float y_norm = ((float)row + 0.5f) / (float)H; - float pe_vec[256]; - sam3_pe_encode_coord(pe_vec, x_norm, y_norm, - state.pe_gauss_cache.data(), num_pos_feats); - for (int d = 0; d < D; ++d) - state.dense_pe_cache[d + col * D + row * D * H] = pe_vec[d]; - } - } - - state.dense_nomask_cache.resize(D * H * H); - for (int row = 0; row < H; ++row) { - for (int col = 0; col < H; ++col) { - for (int d = 0; d < D; ++d) - state.dense_nomask_cache[d + col * D + row * D * H] = state.no_mask_emb_cache[d]; - } - } - - state.pe_cache_valid = true; - fprintf(stderr, "%s: PE cache populated (%d embeddings, %.1f KB dense grids)\n", - __func__, pe_nel, 2.0f * D * H * H * sizeof(float) / 1024.0f); -} - -// Match the official image predictor prompt ordering: -// box corners (if any) → user points → trailing padding point. -// Even when a box is present, SAM keeps the final padding token because boxes -// are merged into the point stream before calling PromptEncoder(points=..., boxes=None). -static void sam3_collect_pvs_prompt_tokens(const sam3_pvs_params& params, - std::vector& all_coords, - std::vector& all_labels) { - all_coords.clear(); - all_labels.clear(); - - if (params.use_box) { - all_coords.push_back(params.box.x0); - all_coords.push_back(params.box.y0); - all_labels.push_back(2); - - all_coords.push_back(params.box.x1); - all_coords.push_back(params.box.y1); - all_labels.push_back(3); - } - - for (const auto& pt : params.pos_points) { - all_coords.push_back(pt.x); - all_coords.push_back(pt.y); - all_labels.push_back(1); - } - for (const auto& pt : params.neg_points) { - all_coords.push_back(pt.x); - all_coords.push_back(pt.y); - all_labels.push_back(0); - } - - all_coords.push_back(0.0f); - all_coords.push_back(0.0f); - all_labels.push_back(-1); -} - -// Build sparse and dense embeddings from point/box prompts -// sparse_out: [D, N_pts, 1] where N_pts = n_pos + n_neg + (use_box ? 2 : 0) + pad -// dense_out: [D, H, H, 1] (no-mask default or mask downsample) -struct sam3_pe_result { - struct ggml_tensor* sparse; // [D, N_pts, 1] - struct ggml_tensor* dense; // [D, H, H, 1] - struct ggml_tensor* image_pe; // [D, H, H, 1] — dense positional encoding grid - int n_tokens; -}; - -static sam3_pe_result sam3_build_sam_pe( - struct ggml_context* ctx, - const sam3_pvs_params& params, - int embed_dim, int feat_size) { - const int D = embed_dim; // 256 - const int H = feat_size; // 72 - - int N_pts = (int)(params.pos_points.size() + params.neg_points.size()) + 1; // pad - if (params.use_box) N_pts += 2; - - auto* sparse = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, N_pts, 1); - ggml_set_name(sparse, "sam_pe_sparse"); - ggml_set_input(sparse); - - auto* image_pe = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, D, H, H, 1); - ggml_set_name(image_pe, "sam_pe_image_pe"); - ggml_set_input(image_pe); - - auto* dense = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, D, H, H, 1); - ggml_set_name(dense, "sam_pe_dense"); - ggml_set_input(dense); - - sam3_pe_result result; - result.sparse = sparse; - result.dense = dense; - result.image_pe = image_pe; - result.n_tokens = N_pts; - return result; -} - -/***************************************************************************** -** SAM mask decoder — graph building (Phase 6, Step 6.2) -*****************************************************************************/ - -// TwoWayAttentionBlock forward -static void sam3_twoway_block_forward( - struct ggml_context* ctx, - struct ggml_tensor*& queries, // [D, N_q, B] — modified in place - struct ggml_tensor*& keys, // [D, N_kv, B] — modified in place - struct ggml_tensor* query_pe, // [D, N_q, B] - struct ggml_tensor* key_pe, // [D, N_kv, B] - const sam3_twoway_block& blk, - int n_heads, - bool skip_first_layer_pe) { - // 1. Self-attention on queries - if (skip_first_layer_pe) { - // Python: queries = self.self_attn(q=queries, k=queries, v=queries) - // No residual connection when skipping first layer PE - queries = sam3_sam_attention(ctx, queries, queries, queries, blk.self_attn, n_heads); - } else { - auto* q = ggml_add(ctx, queries, query_pe); - auto* attn_out = sam3_sam_attention(ctx, q, q, queries, blk.self_attn, n_heads); - queries = ggml_add(ctx, queries, attn_out); - } - queries = sam3_layer_norm(ctx, queries, blk.norm1_w, blk.norm1_b); - if (skip_first_layer_pe) { - ggml_set_name(queries, "dbg_twoway_skip_sa_norm"); - ggml_set_output(queries); - } - - // 2. Cross-attention: tokens attending to image - { - auto* q = ggml_add(ctx, queries, query_pe); - auto* k = ggml_add(ctx, keys, key_pe); - auto* attn_out = sam3_sam_attention(ctx, q, k, keys, blk.ca_tok2img, n_heads); - queries = ggml_add(ctx, queries, attn_out); - queries = sam3_layer_norm(ctx, queries, blk.norm2_w, blk.norm2_b); - } - if (skip_first_layer_pe) { - ggml_set_name(queries, "dbg_twoway_skip_ca_tok2img"); - ggml_set_output(queries); - } - - // 3. MLP on queries (ReLU activation) - { - auto* mlp = ggml_mul_mat(ctx, blk.mlp_fc1_w, queries); - mlp = ggml_add(ctx, mlp, blk.mlp_fc1_b); - mlp = ggml_relu(ctx, mlp); - mlp = ggml_mul_mat(ctx, blk.mlp_fc2_w, mlp); - mlp = ggml_add(ctx, mlp, blk.mlp_fc2_b); - queries = ggml_add(ctx, queries, mlp); - queries = sam3_layer_norm(ctx, queries, blk.norm3_w, blk.norm3_b); - } - if (skip_first_layer_pe) { - ggml_set_name(queries, "dbg_twoway_skip_mlp"); - ggml_set_output(queries); - } - - // 4. Cross-attention: image attending to tokens - { - auto* q = ggml_add(ctx, queries, query_pe); - auto* k = ggml_add(ctx, keys, key_pe); - // Note: q and k are swapped — image (k) attends to tokens (q) - auto* attn_out = sam3_sam_attention(ctx, k, q, queries, blk.ca_img2tok, n_heads); - keys = ggml_add(ctx, keys, attn_out); - keys = sam3_layer_norm(ctx, keys, blk.norm4_w, blk.norm4_b); - } - if (skip_first_layer_pe) { - ggml_set_name(keys, "dbg_twoway_skip_img2tok"); - ggml_set_output(keys); - } -} - -// MLP forward: N layers with ReLU (except last), optional sigmoid on last -static struct ggml_tensor* sam3_mlp_forward( - struct ggml_context* ctx, - struct ggml_tensor* x, - struct ggml_tensor* const* weights, - struct ggml_tensor* const* biases, - int n_layers, - bool sigmoid_output = false) { - for (int i = 0; i < n_layers; ++i) { - x = ggml_mul_mat(ctx, weights[i], x); - x = ggml_add(ctx, x, biases[i]); - if (i < n_layers - 1) { - x = ggml_relu(ctx, x); - } - } - if (sigmoid_output) { - x = ggml_sigmoid(ctx, x); - } - return x; -} - -// Full SAM mask decoder graph -// Inputs: -// image_feats: [D, H, H, 1] — tracker neck features (scale 2 = 72×72) -// image_pe: [D, H, H, 1] — dense positional encoding -// sparse_emb: [D, N_pts, 1] — sparse prompt embeddings -// dense_emb: [D, H, H, 1] — dense prompt embeddings (no_mask default) -// feat_s0: [D, H0, H0, 1] — high-res features (scale 0 = 288×288) -// feat_s1: [D, H1, H1, 1] — mid-res features (scale 1 = 144×144) -// Outputs: sam3_dec_result with masks, iou_pred, obj_score, sam_token_out -struct sam3_dec_result { - struct ggml_tensor* masks; // [288*288, N_masks, 1] - struct ggml_tensor* iou_pred; // [N_masks, 1] - struct ggml_tensor* obj_score; // [1, 1] - struct ggml_tensor* sam_token; // [D, 1] — for object pointer - struct ggml_tensor* mask_tokens; // [D, N_masks, 1] — raw SAM mask tokens -}; - -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); - } - - // Final attention: tokens → image - { - auto* q = ggml_add(ctx, queries, query_pe); - auto* k = ggml_add(ctx, keys, key_pe); - auto* attn_out = sam3_sam_attention(ctx, q, k, keys, dec.final_attn, n_heads); - queries = ggml_add(ctx, queries, attn_out); - queries = sam3_layer_norm(ctx, queries, dec.final_norm_w, dec.final_norm_b); - ggml_set_name(queries, "sam_dec_final_queries"); - } - - // Debug: mark transformer outputs - ggml_set_name(queries, "dbg_dec_queries_out"); - ggml_set_output(queries); - ggml_set_name(keys, "dbg_dec_keys_out"); - ggml_set_output(keys); - - // ── Extract output tokens ──────────────────────────────────────────── - // With pred_obj_scores=True (6 tokens): obj(0), iou(1), masks(2..5) - // With pred_obj_scores=False (5 tokens): iou(0), masks(1..4) - const int s = has_obj_score ? 1 : 0; - auto* iou_token_out = ggml_view_3d(ctx, queries, D, 1, 1, - queries->nb[1], queries->nb[2], - s * queries->nb[1]); - iou_token_out = ggml_cont(ctx, iou_token_out); // [D, 1, 1] - - auto* mask_tokens_out = ggml_view_3d(ctx, queries, D, num_mask_tokens, 1, - queries->nb[1], queries->nb[2], - (s + 1) * queries->nb[1]); - mask_tokens_out = ggml_cont(ctx, mask_tokens_out); // [D, 4, 1] - ggml_set_name(mask_tokens_out, "sam_dec_mask_tokens"); - - struct ggml_tensor* obj_in = nullptr; - if (has_obj_score) { - obj_in = ggml_view_3d(ctx, queries, D, 1, 1, - queries->nb[1], queries->nb[2], 0); - obj_in = ggml_cont(ctx, obj_in); // [D, 1, 1] - } - - // SAM output token = first mask token, used for object pointer - auto* sam_token = ggml_view_2d(ctx, queries, D, 1, - queries->nb[1], (s + 1) * queries->nb[1]); - sam_token = ggml_cont(ctx, sam_token); // [D, 1] - ggml_set_name(sam_token, "sam_dec_sam_token"); - - // Upscale: [D, H*H, 1] → ConvTranspose → high-res masks - auto* src_img = ggml_reshape_4d(ctx, keys, D, H, H, 1); - src_img = ggml_cont(ctx, ggml_permute(ctx, src_img, 2, 0, 1, 3)); - - auto* up1 = ggml_conv_transpose_2d_p0(ctx, sam3_conv_transpose_weight(ctx, dec.up1_w), src_img, 2); - up1 = ggml_add(ctx, up1, ggml_reshape_4d(ctx, dec.up1_b, 1, 1, ggml_nelements(dec.up1_b), 1)); - - auto* fs1 = ggml_cont(ctx, ggml_permute(ctx, feat_s1, 2, 0, 1, 3)); - auto* hs1 = ggml_conv_2d_sk_p0(ctx, dec.conv_s1_w, fs1); // [W, H, 64, B] - hs1 = ggml_add(ctx, hs1, ggml_reshape_4d(ctx, dec.conv_s1_b, 1, 1, 64, 1)); - ggml_set_name(hs1, "sam_dec_feat_s1_proj"); - - // Python: act1(ln1(dc1(src) + feat_s1)) — add before LayerNorm - up1 = ggml_add(ctx, up1, hs1); - up1 = ggml_cont(ctx, ggml_permute(ctx, up1, 1, 2, 0, 3)); - up1 = sam3_layer_norm_2d(ctx, up1, dec.up1_norm_w, dec.up1_norm_b); - - up1 = ggml_gelu_erf(ctx, up1); - - // Permute back to [W, H, C, B] for next deconv - up1 = ggml_cont(ctx, ggml_permute(ctx, up1, 2, 0, 1, 3)); // [144, 144, 64, 1] - - // dc2: ConvTranspose2d(64, 32, k=2, s=2) → [288, 288, 32, 1] - auto* up2 = ggml_conv_transpose_2d_p0(ctx, sam3_conv_transpose_weight(ctx, dec.up2_w), up1, 2); - up2 = ggml_add(ctx, up2, ggml_reshape_4d(ctx, dec.up2_b, 1, 1, ggml_nelements(dec.up2_b), 1)); - - // conv_s0: 1x1 conv on feat_s0 (256→32). feat_s0 is [C, W, H, B] — permute for conv. - auto* fs0 = ggml_cont(ctx, ggml_permute(ctx, feat_s0, 2, 0, 1, 3)); // [W, H, C, B] - auto* hs0 = ggml_conv_2d_sk_p0(ctx, dec.conv_s0_w, fs0); // [W, H, 32, B] - hs0 = ggml_add(ctx, hs0, ggml_reshape_4d(ctx, dec.conv_s0_b, 1, 1, 32, 1)); - ggml_set_name(hs0, "sam_dec_feat_s0_proj"); - - // Python: act2(dc2(upscaled_embedding) + feat_s0) — no LayerNorm here - up2 = ggml_add(ctx, up2, hs0); // both [W, H, 32, B] - - // Permute to [C, W, H, B] for subsequent operations - up2 = ggml_cont(ctx, ggml_permute(ctx, up2, 1, 2, 0, 3)); // [32, 288, 288, 1] - - // GELU activation (exact, matching Python nn.GELU) - up2 = ggml_gelu_erf(ctx, up2); - - // up2: [32, 288, 288, 1] — this is our upscaled_embedding - ggml_set_name(up2, "sam_dec_upscaled"); - - // ── Hypernetwork: predict masks ────────────────────────────────────── - // For each mask token i, pass through 3-layer MLP to get [32] vector - // Then dot product with upscaled_embedding [32, (H*4)^2] to get mask - const int H4 = H * 4; - auto* up_flat = ggml_reshape_3d(ctx, up2, 32, H4 * H4, 1); - - // Process each mask token through its hypernetwork MLP - // mask_tokens_out: [D, 4, 1] - struct ggml_tensor* mask_list[4]; - for (int m = 0; m < num_mask_tokens; ++m) { - // Extract token m: [D, 1, 1] - auto* tok = ggml_view_3d(ctx, mask_tokens_out, D, 1, 1, - mask_tokens_out->nb[1], mask_tokens_out->nb[2], - m * mask_tokens_out->nb[1]); - tok = ggml_cont(ctx, tok); // [D, 1, 1] - - // MLP: 3 layers, 256→256→256→32, ReLU on first two - auto* hyper = sam3_mlp_forward(ctx, tok, - dec.hyper_w[m], dec.hyper_b[m], 3); - // hyper: [32, 1, 1] - - // Dot product: hyper^T @ up_flat → [1, 288*288, 1] - // Use mul_mat: up_flat^T [288*288, 32] @ hyper [32, 1] → [288*288, 1, 1] - auto* mask = ggml_mul_mat(ctx, up_flat, hyper); // [288*288, 1, 1] - mask_list[m] = mask; - } - - // Stack masks: [288*288, 4, 1] - auto* masks = mask_list[0]; - for (int m = 1; m < num_mask_tokens; ++m) { - masks = ggml_concat(ctx, masks, mask_list[m], 1); - } - ggml_set_name(masks, "sam_dec_masks"); - - // ── IoU prediction ─────────────────────────────────────────────────── - // iou_token_out: [D, 1, 1] - auto* iou_pred = sam3_mlp_forward(ctx, iou_token_out, - dec.iou_head_w, dec.iou_head_b, 3, - /*sigmoid_output=*/true); - // iou_pred: [4, 1, 1] → reshape to [4, 1] - iou_pred = ggml_reshape_2d(ctx, iou_pred, num_mask_tokens, 1); - ggml_set_name(iou_pred, "sam_dec_iou"); - - // ── Object score ───────────────────────────────────────────────────── - struct ggml_tensor* obj_score; - if (has_obj_score) { - // obj_in: [D, 1, 1] → MLP → [1, 1, 1] - obj_score = sam3_mlp_forward(ctx, obj_in, - dec.obj_head_w, dec.obj_head_b, 3); - obj_score = ggml_reshape_2d(ctx, obj_score, 1, 1); - } else { - // No obj_score prediction — return raw logit 10.0 (sigmoid ≈ 1.0, object always present) - obj_score = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, 1); - ggml_set_name(obj_score, "sam_dec_obj_score"); - ggml_set_input(obj_score); - // Mark that callers must set this to 10.0f before compute - } - - sam3_dec_result res; - res.masks = masks; - res.iou_pred = iou_pred; - res.obj_score = obj_score; - res.sam_token = sam_token; - res.mask_tokens = mask_tokens_out; - return res; -} - -/***************************************************************************** -** Image segmentation — PVS (visual-prompted) (Phase 6, Step 6.3) -*****************************************************************************/ - -sam3_result sam3_segment_pvs(sam3_state& state, - const sam3_model& model, - const sam3_pvs_params& params) { -#if SAM3_LOG_LEVEL >= 1 - auto t_start = std::chrono::high_resolution_clock::now(); -#endif - const auto& hp = model.hparams; - const int D = hp.sam_embed_dim; // 256 - const int H = sam3_eff_feat_size(state, hp); - const int num_mask_tokens = hp.sam_n_multimask + 1; // 4 - const int eff_img_size = sam3_eff_img_size(state, hp); - sam3_result result; - - // ── Validate ───────────────────────────────────────────────────────── - if (!state.neck_trk[0]) { - fprintf(stderr, "%s: image not encoded — call sam3_encode_image first\n", __func__); - return result; - } - if (params.pos_points.empty() && !params.use_box) { - fprintf(stderr, "%s: no prompts provided (need at least one point or box)\n", __func__); - return result; - } - - SAM3_LOG(2, "%s: %zu pos points, %zu neg points, box=%s, multimask=%s\n", - __func__, params.pos_points.size(), params.neg_points.size(), - params.use_box ? "yes" : "no", params.multimask ? "yes" : "no"); - - // ── Build computation graph ────────────────────────────────────────── - const size_t buf_size = ggml_tensor_overhead() * 8192 + ggml_graph_overhead() * 2; - struct ggml_init_params gparams = { - /*.mem_size =*/buf_size, - /*.mem_buffer =*/nullptr, - /*.no_alloc =*/true, - }; - struct ggml_context* ctx0 = ggml_init(gparams); - if (!ctx0) { - fprintf(stderr, "%s: failed to init compute context\n", __func__); - return result; - } - - // ── SAM prompt encoder (CPU pre-compute + input tensors) ───────────── - auto pe_out = sam3_build_sam_pe(ctx0, params, D, H); - - // ── Create fresh input tensors for tracker features ────────────────── - // CRITICAL: Do NOT use state.neck_trk[*] directly in graph operations - // (like ggml_add). They are tensors from a previous graph, and - // ggml_build_forward_expand traces all ancestors — which would pull in - // the ENTIRE ViT + neck recomputation (2500+ nodes, ~37 seconds). - // Instead: create fresh input tensors and copy data from state on CPU. - const int H0 = H * 4; // 288 - const int H1 = H * 2; // 144 - - auto* image_feats = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, D, H, H, 1); - ggml_set_name(image_feats, "sam_dec_image_feats"); - ggml_set_input(image_feats); - - auto* feat_s0 = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, D, H0, H0, 1); - ggml_set_name(feat_s0, "pvs_feat_s0"); - ggml_set_input(feat_s0); - - auto* feat_s1 = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, D, H1, H1, 1); - ggml_set_name(feat_s1, "pvs_feat_s1"); - ggml_set_input(feat_s1); - - // ── SAM mask decoder graph ─────────────────────────────────────────── - auto dec_out = sam3_build_sam_dec_graph(ctx0, model, - image_feats, - pe_out.image_pe, - pe_out.sparse, - pe_out.dense, - feat_s0, - feat_s1, H); - - // Mark outputs - ggml_set_output(dec_out.masks); - ggml_set_output(dec_out.iou_pred); - ggml_set_output(dec_out.obj_score); - ggml_set_output(dec_out.sam_token); - - // ── Build and allocate graph ───────────────────────────────────────── - struct ggml_cgraph* graph = ggml_new_graph_custom(ctx0, 32768, false); - ggml_build_forward_expand(graph, dec_out.masks); - ggml_build_forward_expand(graph, dec_out.iou_pred); - ggml_build_forward_expand(graph, dec_out.obj_score); - ggml_build_forward_expand(graph, dec_out.sam_token); - - auto* galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); - if (!ggml_gallocr_reserve(galloc, graph)) { - fprintf(stderr, "%s: failed to reserve graph memory\n", __func__); - ggml_gallocr_free(galloc); - ggml_free(ctx0); - return result; - } - if (!ggml_gallocr_alloc_graph(galloc, graph)) { - fprintf(stderr, "%s: failed to allocate graph\n", __func__); - ggml_gallocr_free(galloc); - ggml_free(ctx0); - return result; - } - - SAM3_LOG(2, "%s: graph allocated, %d nodes\n", __func__, ggml_graph_n_nodes(graph)); - - // Set default obj_score when pred_obj_scores=False (older SAM2 models) - if (!model.sam_dec.obj_score_token) { - auto* t = ggml_graph_get_tensor(graph, "sam_dec_obj_score"); - if (t) { float v = 10.0f; ggml_backend_tensor_set(t, &v, 0, sizeof(float)); } - } - - // ── Upload input data (using cached embeddings) ──────────────────── - // Populate PE cache on first call (reads model weights from GPU once) - sam3_populate_pe_cache(state, model); - - { - const int N_pts = pe_out.n_tokens; - const int num_pos_feats = D / 2; - - std::vector all_coords; - std::vector all_labels; - sam3_collect_pvs_prompt_tokens(params, all_coords, all_labels); - - // Sparse embeddings — only this changes per call - std::vector sparse_data(N_pts * D, 0.0f); - for (int p = 0; p < N_pts; ++p) { - // Scale from original image space to model input space, then shift to pixel center - float px = all_coords[p * 2 + 0] / (float)state.orig_width * (float)eff_img_size + 0.5f; - float py = all_coords[p * 2 + 1] / (float)state.orig_height * (float)eff_img_size + 0.5f; - float x_norm = px / (float)eff_img_size; - float y_norm = py / (float)eff_img_size; - float pe_vec[256]; - sam3_pe_encode_coord(pe_vec, x_norm, y_norm, - state.pe_gauss_cache.data(), num_pos_feats); - int label = all_labels[p]; - if (label == -1) { - for (int d = 0; d < D; ++d) - sparse_data[p * D + d] = state.not_a_point_cache[d]; - } else { - for (int d = 0; d < D; ++d) - sparse_data[p * D + d] = pe_vec[d] + state.point_emb_cache[label][d]; - } - } - ggml_backend_tensor_set(pe_out.sparse, sparse_data.data(), 0, N_pts * D * sizeof(float)); - - // Dump sparse embeddings if requested - { - const char* dd = getenv("SAM2_DUMP_DIR"); - if (dd) { - char p[512]; snprintf(p, sizeof(p), "%s/cpp_sparse_emb.bin", dd); - FILE* f = fopen(p, "wb"); - if (f) { fwrite(sparse_data.data(), sizeof(float), N_pts * D, f); fclose(f); } - fprintf(stderr, " [DUMP] cpp_sparse_emb: %d tokens x %d dims\n", N_pts, D); - } - } - - // Dense PE grid and no-mask embedding — use pre-computed caches - ggml_backend_tensor_set(pe_out.image_pe, state.dense_pe_cache.data(), - 0, D * H * H * sizeof(float)); - ggml_backend_tensor_set(pe_out.dense, state.dense_nomask_cache.data(), - 0, D * H * H * sizeof(float)); - } - - // ── Copy tracker features from state to fresh input tensors ───────── - // image_feats = neck_trk[2] + no_mem_embed (computed on CPU) - { - const int n2 = D * H * H; - std::vector trk2(n2), no_mem_data(D); - ggml_backend_tensor_get(state.neck_trk[2], trk2.data(), 0, n2 * sizeof(float)); - ggml_backend_tensor_get(model.tensors.at("no_mem_embed"), no_mem_data.data(), 0, D * sizeof(float)); - // Add no_mem_embed (broadcast [D] to [D, H, H]) - for (int s = 0; s < H * H; ++s) - for (int d = 0; d < D; ++d) - trk2[d + s * D] += no_mem_data[d]; - ggml_backend_tensor_set(image_feats, trk2.data(), 0, n2 * sizeof(float)); - - // Dump image_feats (with no_mem_embed) if requested - { - const char* dd = getenv("SAM2_DUMP_DIR"); - if (dd) { - char p[512]; snprintf(p, sizeof(p), "%s/cpp_image_feats.bin", dd); - FILE* f = fopen(p, "wb"); - if (f) { fwrite(trk2.data(), sizeof(float), n2, f); fclose(f); } - fprintf(stderr, " [DUMP] cpp_image_feats: [%d, %d, %d]\n", D, H, H); - } - } - - // feat_s0 = neck_trk[0], feat_s1 = neck_trk[1] - const int n0 = D * H0 * H0; - const int n1 = D * H1 * H1; - std::vector trk0(n0), trk1(n1); - ggml_backend_tensor_get(state.neck_trk[0], trk0.data(), 0, n0 * sizeof(float)); - ggml_backend_tensor_set(feat_s0, trk0.data(), 0, n0 * sizeof(float)); - ggml_backend_tensor_get(state.neck_trk[1], trk1.data(), 0, n1 * sizeof(float)); - ggml_backend_tensor_set(feat_s1, trk1.data(), 0, n1 * sizeof(float)); - } - - // ── Compute ────────────────────────────────────────────────────────── - { -#if SAM3_LOG_LEVEL >= 1 - auto t0 = std::chrono::high_resolution_clock::now(); -#endif - if (!sam3_graph_compute(model.backend, graph, state.n_threads)) { - ggml_gallocr_free(galloc); - ggml_free(ctx0); - return result; - } -#if SAM3_LOG_LEVEL >= 1 - auto t1 = std::chrono::high_resolution_clock::now(); - double ms = std::chrono::duration(t1 - t0).count(); - SAM3_LOG(1, "%s: graph computed in %.1f ms (%d threads)\n", - __func__, ms, state.n_threads); -#endif - } - - // ── Dump decoder outputs if SAM2_DUMP_DIR set ────────────────────── - { - const char* dump_dir = getenv("SAM2_DUMP_DIR"); - if (dump_dir) { - auto dump_t = [&](const char* name, struct ggml_tensor* t) { - if (!t) return; - int64_t nb = ggml_nbytes(t); - std::vector buf(nb); - ggml_backend_tensor_get(t, buf.data(), 0, nb); - char path[512]; - snprintf(path, sizeof(path), "%s/%s.bin", dump_dir, name); - FILE* f = fopen(path, "wb"); - if (f) { fwrite(buf.data(), 1, nb, f); fclose(f); } - snprintf(path, sizeof(path), "%s/%s.shape", dump_dir, name); - f = fopen(path, "w"); - if (f) { - fprintf(f, "%lld,%lld,%lld,%lld", - (long long)t->ne[0], (long long)t->ne[1], - (long long)t->ne[2], (long long)t->ne[3]); - fclose(f); - } - fprintf(stderr, " [DUMP] %s: [%lld,%lld,%lld,%lld]\n", name, - (long long)t->ne[0], (long long)t->ne[1], - (long long)t->ne[2], (long long)t->ne[3]); - }; - dump_t("cpp_pvs_masks", dec_out.masks); - dump_t("cpp_pvs_iou", dec_out.iou_pred); - dump_t("cpp_pvs_obj_score", dec_out.obj_score); - // Decoder transformer intermediates - dump_t("cpp_dec_queries", ggml_graph_get_tensor(graph, "dbg_dec_queries_out")); - dump_t("cpp_dec_keys", ggml_graph_get_tensor(graph, "dbg_dec_keys_out")); - // Block 0 internals - const char* b0_names[] = {"sam_dec_tokens_initial", - "dbg_twoway_skip_sa_norm", "dbg_twoway_skip_ca_tok2img", - "dbg_twoway_skip_mlp", "dbg_twoway_skip_img2tok", - "dbg_sa0_Q_proj", "dbg_sa0_merged"}; - for (auto* bn : b0_names) dump_t(bn, ggml_graph_get_tensor(graph, bn)); - // Per-block outputs - for (int bi = 0; bi < 2; bi++) { - char bn[64]; - snprintf(bn, sizeof(bn), "sam_dec_block%d_queries", bi); - dump_t(bn, ggml_graph_get_tensor(graph, bn)); - snprintf(bn, sizeof(bn), "sam_dec_block%d_keys", bi); - dump_t(bn, ggml_graph_get_tensor(graph, bn)); - } - } - } - - // ── Read outputs ───────────────────────────────────────────────────── - // masks: [H*4×H*4, 4, 1] - const int mask_hw = H * 4; - std::vector masks_data(mask_hw * mask_hw * num_mask_tokens); - ggml_backend_tensor_get(dec_out.masks, masks_data.data(), 0, masks_data.size() * sizeof(float)); - - // IoU predictions: [4, 1] - std::vector iou_data(num_mask_tokens); - ggml_backend_tensor_get(dec_out.iou_pred, iou_data.data(), 0, num_mask_tokens * sizeof(float)); - - // Object score: [1, 1] - float obj_logit = 0.0f; - ggml_backend_tensor_get(dec_out.obj_score, &obj_logit, 0, sizeof(float)); - float obj_score = 1.0f / (1.0f + expf(-obj_logit)); - - // SAM output token: [D, 1] — needed for object pointer extraction in tracking - std::vector sam_token_data(D); - ggml_backend_tensor_get(dec_out.sam_token, sam_token_data.data(), 0, D * sizeof(float)); - - SAM3_LOG(2, "%s: obj_score=%.4f (logit=%.4f), iou=[%.3f, %.3f, %.3f, %.3f]\n", - __func__, obj_score, obj_logit, - iou_data[0], iou_data[1], iou_data[2], iou_data[3]); - - // ── Select masks based on multimask mode ───────────────────────────── - // Python: if multimask_output → masks[:, 1:, :, :], iou_pred[:, 1:] - // else → masks[:, 0:1, :, :], iou_pred[:, 0:1] - int start_idx, end_idx; - if (params.multimask) { - start_idx = 1; - end_idx = num_mask_tokens; - } else { - start_idx = 0; - end_idx = 1; - } - - for (int m = start_idx; m < end_idx; ++m) { - sam3_detection det; - det.sam_token = sam_token_data; - - // Resize mask from 288×288 to original image size - const float* mask_ptr = masks_data.data() + m * mask_hw * mask_hw; - auto mask_resized = sam3_bilinear_interpolate(mask_ptr, mask_hw, mask_hw, - state.orig_width, state.orig_height); - - // Binarize at threshold 0.0 (sigmoid(logit) > 0.5 ↔ logit > 0.0) - det.mask.width = state.orig_width; - det.mask.height = state.orig_height; - det.mask.data.resize(state.orig_width * state.orig_height); - for (int i = 0; i < (int)mask_resized.size(); ++i) { - det.mask.data[i] = (mask_resized[i] > 0.0f) ? 255 : 0; - } - - det.mask.iou_score = iou_data[m]; - det.mask.obj_score = obj_score; - det.mask.instance_id = m; - det.score = iou_data[m]; - det.iou_score = iou_data[m]; - det.instance_id = m; - - // Compute bounding box from mask - int min_x = state.orig_width, min_y = state.orig_height; - int max_x = 0, max_y = 0; - for (int y = 0; y < state.orig_height; ++y) { - for (int x = 0; x < state.orig_width; ++x) { - if (det.mask.data[y * state.orig_width + x] > 0) { - min_x = std::min(min_x, x); - min_y = std::min(min_y, y); - max_x = std::max(max_x, x); - max_y = std::max(max_y, y); - } - } - } - det.box = {(float)min_x, (float)min_y, (float)max_x, (float)max_y}; - - result.detections.push_back(std::move(det)); - } - - SAM3_LOG(2, "%s: %zu masks returned\n", __func__, result.detections.size()); - - // ── Cleanup ────────────────────────────────────────────────────────── - ggml_gallocr_free(galloc); - ggml_free(ctx0); - -#if SAM3_LOG_LEVEL >= 1 - auto t_end = std::chrono::high_resolution_clock::now(); - double total_ms = std::chrono::duration(t_end - t_start).count(); - SAM3_LOG(1, "%s: completed in %.1f ms\n", __func__, total_ms); -#endif - - return result; -} - -/***************************************************************************** -** Video tracking (Phase 7) -*****************************************************************************/ - -struct sam3_prop_output { - std::vector mask_logits; - std::vector iou_scores; - float obj_score; - std::vector sam_token; - int n_masks, mask_h, mask_w; -}; - -// Lazily compute and cache PE/RoPE data that is identical across all propagation calls. -static void sam3_ensure_tracker_pe_caches(sam3_tracker& tracker, const sam3_hparams& hp, - int eff_feat_size = 0) { - const int H = (eff_feat_size > 0) ? eff_feat_size : hp.feat_size(); - if (tracker.pe_caches_valid) return; - - const int D = hp.neck_dim; // 256 - const int MD = hp.mem_out_dim; // 64 - const int N = H * H; - const int half_d = D / 2; // 128 - - tracker.cached_sinpe_256 = sam3_sinusoidal_pe_2d(H, H, D); - tracker.cached_sinpe_64 = sam3_sinusoidal_pe_2d(H, H, MD); - - // Compute axial CIS and reorder to [2, half_d, N] layout - std::vector rope_raw(N * D); - sam3_compute_axial_cis(rope_raw.data(), D, H, H, 10000.0f, 1.0f); - tracker.cached_axial_cis_reord.resize(2 * half_d * N); - for (int n = 0; n < N; ++n) - for (int i = 0; i < half_d; ++i) { - tracker.cached_axial_cis_reord[0 + i * 2 + n * D] = rope_raw[n * D + i * 2 + 0]; - tracker.cached_axial_cis_reord[1 + i * 2 + n * D] = rope_raw[n * D + i * 2 + 1]; - } - - // EdgeTAM: compute 16x16 RoPE for cross-attention K (perceiver 2D latents) - if (hp.has_perceiver) { - const int K16 = 16; // perceiver 2D grid size - const int NK = K16 * K16; // 256 tokens - std::vector rope_k16_raw(NK * D); - sam3_compute_axial_cis(rope_k16_raw.data(), D, K16, K16, 10000.0f, 1.0f); - tracker.cached_axial_cis_k16_reord.resize(2 * half_d * NK); - for (int n = 0; n < NK; ++n) - for (int i = 0; i < half_d; ++i) { - tracker.cached_axial_cis_k16_reord[0 + i * 2 + n * D] = rope_k16_raw[n * D + i * 2 + 0]; - tracker.cached_axial_cis_k16_reord[1 + i * 2 + n * D] = rope_k16_raw[n * D + i * 2 + 1]; - } - SAM3_LOG(2, "%s: EdgeTAM K16 RoPE cache: %zu floats\n", __func__, - tracker.cached_axial_cis_k16_reord.size()); - } - - tracker.pe_caches_valid = true; - SAM3_LOG(2, "%s: tracker PE caches populated (%.1f KB)\n", __func__, - (tracker.cached_sinpe_256.size() + tracker.cached_sinpe_64.size() + - tracker.cached_axial_cis_reord.size()) * sizeof(float) / 1024.0f); -} - -static sam3_prop_output sam3_propagate_single( - sam3_tracker& tracker, sam3_state& state, const sam3_model& model, - const sam3_masklet& masklet, - const std::vector& mem_bank, - const std::vector>& ptr_bank) { - sam3_prop_output output = {}; - const auto& hp = model.hparams; - const int D = hp.neck_dim, MD = hp.mem_out_dim; - const int H = sam3_eff_feat_size(state, hp); - const int N = H * H; - - auto sel = sam3_select_memory_frames(mem_bank, hp.num_maskmem); - if (sel.empty()) return output; - - // ── Build prompt and prompt_pos via sam3_build_prompt_and_pos ───────── - // EdgeTAM with perceiver: each slot has 512 tokens (not H*H). - // Standard SAM2/SAM3: each slot has H*H tokens. - const bool use_perceiver = hp.has_perceiver != 0; - const int N_per_slot = use_perceiver ? (hp.perceiver_n_latents_1d + hp.perceiver_n_latents_2d) : N; - - int n_sel = (int)sel.size(); - std::vector> slot_feats(n_sel), slot_pes(n_sel); - std::vector spatial_tpos(n_sel, 1); // default t_pos=1 for non-cond - for (int s = 0; s < n_sel; ++s) { - slot_feats[s].resize(MD * N_per_slot); - ggml_backend_tensor_get(mem_bank[sel[s]].spatial_feats, - slot_feats[s].data(), 0, MD * N_per_slot * sizeof(float)); - slot_pes[s].resize(MD * N_per_slot); - if (mem_bank[sel[s]].spatial_pe) { - ggml_backend_tensor_get(mem_bank[sel[s]].spatial_pe, - slot_pes[s].data(), 0, MD * N_per_slot * sizeof(float)); - } else { - sam3_ensure_tracker_pe_caches(tracker, hp, H); - if (use_perceiver) { - // For perceiver: zeros for 1D tokens, sinusoidal for 2D tokens - const int N_1d = hp.perceiver_n_latents_1d; - const int N_2d = hp.perceiver_n_latents_2d; - memset(slot_pes[s].data(), 0, MD * N_1d * sizeof(float)); - auto pe_2d = sam3_sinusoidal_pe_2d(16, 16, MD); - memcpy(slot_pes[s].data() + MD * N_1d, pe_2d.data(), MD * N_2d * sizeof(float)); - } else { - slot_pes[s] = tracker.cached_sinpe_64; - } - } - spatial_tpos[s] = mem_bank[sel[s]].is_cond_frame ? 0 : (n_sel - s); - } - - int P = std::min((int)ptr_bank.size(), hp.max_obj_ptrs); - std::vector> obj_ptrs(P); - std::vector ptr_tpos(P); - int cur_frame = tracker.frame_index; - for (int p = 0; p < P; ++p) { - obj_ptrs[p].resize(D); - ggml_backend_tensor_get(ptr_bank[p].second, obj_ptrs[p].data(), 0, D * sizeof(float)); - // Use actual frame distance (matches Python: abs(frame_idx - t)) - ptr_tpos[p] = std::abs(cur_frame - ptr_bank[p].first); - if (ptr_tpos[p] < 1) ptr_tpos[p] = 1; // minimum distance of 1 - } - - auto pd = sam3_build_prompt_and_pos(model, slot_feats, slot_pes, spatial_tpos, obj_ptrs, ptr_tpos, H); - - // ── RoPE frequencies (cached) ────────────────────────────────────── - sam3_ensure_tracker_pe_caches(tracker, hp, H); - const int half_d = D / 2; // 128 - const auto& rope_q_reord = tracker.cached_axial_cis_reord; - // For cross-attn K: build rope_k_data for all M_spatial tokens - std::vector rope_k_data; - if (pd.M_spatial > 0) { - rope_k_data.resize(2 * half_d * pd.M_spatial); - if (use_perceiver) { - // EdgeTAM perceiver: each frame has N_per_slot=512 tokens. - // First 256 (1D latents): identity RoPE (cos=1, sin=0). - // Last 256 (2D latents): 16x16 RoPE from cached_axial_cis_k16_reord. - const int N_1d = hp.perceiver_n_latents_1d; // 256 - const int N_2d = hp.perceiver_n_latents_2d; // 256 - const auto& rope_k16 = tracker.cached_axial_cis_k16_reord; // [2, 128, 256] - for (int s = 0; s < n_sel; ++s) { - float* dst = rope_k_data.data() + s * D * N_per_slot; - // 1D tokens: identity RoPE (cos=1, sin=0) in [2, half_d, N_1d] layout - // Layout: for token n, dim i: cos at [0 + i*2 + n*D], sin at [1 + i*2 + n*D] - for (int n = 0; n < N_1d; ++n) - for (int i = 0; i < half_d; ++i) { - dst[0 + i * 2 + n * D] = 1.0f; // cos = 1 - dst[1 + i * 2 + n * D] = 0.0f; // sin = 0 - } - // 2D tokens: copy 16x16 RoPE - float* dst_2d = dst + D * N_1d; - memcpy(dst_2d, rope_k16.data(), D * N_2d * sizeof(float)); - } - } else { - // Standard: repeat HxH axial CIS for each memory frame - for (int s = 0; s < pd.M_spatial / N; ++s) - memcpy(rope_k_data.data() + s * D * N, rope_q_reord.data(), D * N * sizeof(float)); - } - } - - // ── Build graph ───────────────────────────────────────────────────── - const size_t buf_size = ggml_tensor_overhead() * 32768 + ggml_graph_overhead() * 2; - struct ggml_init_params gparams = {buf_size, nullptr, true}; - auto* ctx0 = ggml_init(gparams); - if (!ctx0) return output; - - // CRITICAL: create fresh input tensors for state features. - // Using state.neck_trk[*] directly as ggml_reshape operands pulls in - // the entire ViT+neck recomputation from the image encoder graph. - auto* curr = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, D, N, 1); - ggml_set_name(curr, "prop_curr"); - ggml_set_input(curr); - - // src_pos (sinusoidal PE 256-dim for 72×72) - auto* src_pos_t = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, D, N, 1); - ggml_set_name(src_pos_t, "src_pos"); - ggml_set_input(src_pos_t); - - // Prompt and prompt_pos - auto* prompt_t = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, MD, pd.M_total, 1); - ggml_set_name(prompt_t, "prompt"); - ggml_set_input(prompt_t); - auto* prompt_pos_t = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, MD, pd.M_total, 1); - ggml_set_name(prompt_pos_t, "prompt_pos"); - ggml_set_input(prompt_pos_t); - - // RoPE frequencies - auto* rope_q_t = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 2, half_d, N); - ggml_set_name(rope_q_t, "rope_q"); - ggml_set_input(rope_q_t); - struct ggml_tensor* rope_k_t = nullptr; - if (pd.M_spatial > 0) { - rope_k_t = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 2, half_d, pd.M_spatial); - ggml_set_name(rope_k_t, "rope_k"); - ggml_set_input(rope_k_t); - } - - auto* conditioned = sam3_build_mem_attn_graph(ctx0, model, curr, src_pos_t, - prompt_t, prompt_pos_t, - rope_q_t, rope_k_t, - pd.num_obj_ptr_tokens); - auto* cond_spatial = ggml_reshape_4d(ctx0, conditioned, D, H, H, 1); - - // Bug 3 fix: single not_a_point_embed token instead of empty sparse - auto* sparse_in = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, D, 1, 1); - ggml_set_name(sparse_in, "prop_sparse"); - ggml_set_input(sparse_in); - - auto* image_pe = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, D, H, H, 1); - ggml_set_name(image_pe, "prop_pe"); - ggml_set_input(image_pe); - auto* dense_emb = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, D, H, H, 1); - ggml_set_name(dense_emb, "prop_dense"); - ggml_set_input(dense_emb); - - const int H0 = H * 4, H1 = H * 2; - auto* trk_s0 = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, D, H0, H0, 1); - ggml_set_name(trk_s0, "prop_trk_s0"); - ggml_set_input(trk_s0); - auto* trk_s1 = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, D, H1, H1, 1); - ggml_set_name(trk_s1, "prop_trk_s1"); - ggml_set_input(trk_s1); - - auto dec = sam3_build_sam_dec_graph(ctx0, model, cond_spatial, image_pe, - sparse_in, dense_emb, - trk_s0, trk_s1, H); - ggml_set_output(dec.masks); - ggml_set_output(dec.iou_pred); - ggml_set_output(dec.obj_score); - ggml_set_output(dec.sam_token); - if (dec.mask_tokens) ggml_set_output(dec.mask_tokens); - - auto* graph = ggml_new_graph_custom(ctx0, 32768, false); - ggml_build_forward_expand(graph, dec.masks); - ggml_build_forward_expand(graph, dec.iou_pred); - ggml_build_forward_expand(graph, dec.obj_score); - ggml_build_forward_expand(graph, dec.sam_token); - if (dec.mask_tokens) ggml_build_forward_expand(graph, dec.mask_tokens); - - auto* galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); - if (!ggml_gallocr_reserve(galloc, graph) || !ggml_gallocr_alloc_graph(galloc, graph)) { - ggml_gallocr_free(galloc); - ggml_free(ctx0); - return output; - } - - // Upload prompt data - ggml_backend_tensor_set(prompt_t, pd.prompt.data(), 0, pd.prompt.size() * sizeof(float)); - ggml_backend_tensor_set(prompt_pos_t, pd.prompt_pos.data(), 0, pd.prompt_pos.size() * sizeof(float)); - ggml_backend_tensor_set(rope_q_t, rope_q_reord.data(), 0, rope_q_reord.size() * sizeof(float)); - if (rope_k_t && !rope_k_data.empty()) - ggml_backend_tensor_set(rope_k_t, rope_k_data.data(), 0, rope_k_data.size() * sizeof(float)); - - // Set default obj_score when pred_obj_scores=False (older SAM2 models) - if (!model.sam_dec.obj_score_token) { - auto* t = ggml_graph_get_tensor(graph, "sam_dec_obj_score"); - if (t) { float v = 10.0f; ggml_backend_tensor_set(t, &v, 0, sizeof(float)); } - } - - // Upload src_pos (sinusoidal PE 256-dim) - ggml_backend_tensor_set(src_pos_t, tracker.cached_sinpe_256.data(), 0, - tracker.cached_sinpe_256.size() * sizeof(float)); - - // Upload not_a_point_embed, image_pe, dense_emb from state PE cache - sam3_populate_pe_cache(state, model); - ggml_backend_tensor_set(sparse_in, state.not_a_point_cache, 0, D * sizeof(float)); - ggml_backend_tensor_set(image_pe, state.dense_pe_cache.data(), 0, - state.dense_pe_cache.size() * sizeof(float)); - ggml_backend_tensor_set(dense_emb, state.dense_nomask_cache.data(), 0, - state.dense_nomask_cache.size() * sizeof(float)); - - // Copy tracker features from state to fresh input tensors - { - std::vector c2(D * N); - ggml_backend_tensor_get(state.neck_trk[2], c2.data(), 0, D * N * sizeof(float)); - ggml_backend_tensor_set(curr, c2.data(), 0, D * N * sizeof(float)); - - std::vector s0(D * H0 * H0); - ggml_backend_tensor_get(state.neck_trk[0], s0.data(), 0, D * H0 * H0 * sizeof(float)); - ggml_backend_tensor_set(trk_s0, s0.data(), 0, D * H0 * H0 * sizeof(float)); - - std::vector s1(D * H1 * H1); - ggml_backend_tensor_get(state.neck_trk[1], s1.data(), 0, D * H1 * H1 * sizeof(float)); - ggml_backend_tensor_set(trk_s1, s1.data(), 0, D * H1 * H1 * sizeof(float)); - } - - if (!sam3_graph_compute(model.backend, graph, 4)) { - ggml_gallocr_free(galloc); - ggml_free(ctx0); - return output; - } - - const int mhw = H * 4; - const int num_mask_tokens = hp.sam_n_multimask + 1; // 4 - - // SAM2/EdgeTAM with multimask_output_in_sam: read all masks, select best by IoU. - // SAM3 / SAM2 without multimask: use first mask only. - bool use_multimask = (hp.is_sam2() || hp.is_edgetam()) && hp.multimask_output_in_sam; - - if (use_multimask) { - // Read all 4 IoU predictions - std::vector all_ious(num_mask_tokens); - ggml_backend_tensor_get(dec.iou_pred, all_ious.data(), 0, num_mask_tokens * sizeof(float)); - - // Multimask uses mask tokens 1-3 (skip token 0 which is the single-mask output) - int best_idx = 1; - float best_iou = all_ious[1]; - for (int m = 2; m < num_mask_tokens; ++m) { - if (all_ious[m] > best_iou) { best_iou = all_ious[m]; best_idx = m; } - } - - output.n_masks = 1; - output.mask_h = mhw; - output.mask_w = mhw; - output.mask_logits.resize(mhw * mhw); - // Read best mask (offset by best_idx * mhw * mhw) - ggml_backend_tensor_get(dec.masks, output.mask_logits.data(), - best_idx * mhw * mhw * sizeof(float), mhw * mhw * sizeof(float)); - output.iou_scores.resize(1); - output.iou_scores[0] = best_iou; - ggml_backend_tensor_get(dec.obj_score, &output.obj_score, 0, sizeof(float)); - - // Object pointer token: use best multimask token if use_multimask_token_for_obj_ptr - output.sam_token.resize(D); - if (hp.use_multimask_token_for_obj_ptr && dec.mask_tokens) { - ggml_backend_tensor_get(dec.mask_tokens, output.sam_token.data(), - best_idx * D * sizeof(float), D * sizeof(float)); - } else { - ggml_backend_tensor_get(dec.sam_token, output.sam_token.data(), 0, D * sizeof(float)); - } - } else { - output.n_masks = 1; - output.mask_h = mhw; - output.mask_w = mhw; - output.mask_logits.resize(mhw * mhw); - ggml_backend_tensor_get(dec.masks, output.mask_logits.data(), 0, mhw * mhw * sizeof(float)); - output.iou_scores.resize(1); - ggml_backend_tensor_get(dec.iou_pred, output.iou_scores.data(), 0, sizeof(float)); - ggml_backend_tensor_get(dec.obj_score, &output.obj_score, 0, sizeof(float)); - output.sam_token.resize(D); - ggml_backend_tensor_get(dec.sam_token, output.sam_token.data(), 0, D * sizeof(float)); - } - - ggml_gallocr_free(galloc); - ggml_free(ctx0); - return output; -} - -static std::vector> sam3_match_detections( - const std::vector& masklets, const std::vector& dets, - const std::vector& prop_masks, float iou_threshold) { - std::vector> matches; - if (masklets.empty() || dets.empty()) return matches; - int n_m = (int)masklets.size(), n_d = (int)dets.size(); - std::vector dm(n_d, false); - for (int i = 0; i < n_m; ++i) { - if (i >= (int)prop_masks.size() || prop_masks[i].data.empty()) continue; - int bj = -1; - float bi = iou_threshold; - for (int j = 0; j < n_d; ++j) { - if (dm[j] || dets[j].mask.data.empty()) continue; - int w = prop_masks[i].width, h = prop_masks[i].height; - if (w != dets[j].mask.width || h != dets[j].mask.height) continue; - float iou = sam3_mask_iou(prop_masks[i].data.data(), dets[j].mask.data.data(), w * h); - if (iou > bi) { - bi = iou; - bj = j; - } - } - if (bj >= 0) { - matches.push_back({i, bj}); - dm[bj] = true; - } - } - return matches; -} - -static void sam3_update_tracker(sam3_tracker& tracker, int frame_idx) { - for (auto it = tracker.pending.begin(); it != tracker.pending.end();) { - int age = frame_idx - it->first_frame; - if (age >= tracker.params.hotstart_delay && it->mds_sum > 0) { - it->confirmed = true; - tracker.masklets.push_back(std::move(*it)); - it = tracker.pending.erase(it); - } else if (age >= tracker.params.hotstart_delay) { - it = tracker.pending.erase(it); - } else - ++it; - } - for (auto it = tracker.masklets.begin(); it != tracker.masklets.end();) { - if (frame_idx - it->last_seen > tracker.params.max_keep_alive) { - tracker.mem_banks.erase(it->instance_id); - tracker.ptr_banks.erase(it->instance_id); - it = tracker.masklets.erase(it); - } else - ++it; - } -} - -static bool sam3_encode_memory( - sam3_tracker& tracker, sam3_state& state, const sam3_model& model, - int inst_id, const float* mask_logits, int mask_h, int mask_w, - int frame_idx, bool is_cond, float obj_score) { - const auto& hp = model.hparams; - const int D = hp.neck_dim, MD = hp.mem_out_dim; - const int H = sam3_eff_feat_size(state, hp); - const int HIGH_RES = sam3_eff_img_size(state, hp); - const int INTERPOL = H * 16; - - // Mask preprocessing: mask_logits → HIGH_RES → sigmoid → scale/bias → INTERPOL - auto m_hires = sam3_bilinear_interpolate(mask_logits, mask_w, mask_h, HIGH_RES, HIGH_RES); - const float sig_scale = hp.sigmoid_scale(), sig_bias = hp.sigmoid_bias(); - for (auto& v : m_hires) { float s = 1.0f / (1.0f + expf(-v)); v = s * sig_scale + sig_bias; } - auto m_interp = sam3_bilinear_interpolate(m_hires.data(), HIGH_RES, HIGH_RES, INTERPOL, INTERPOL); - - const size_t bs = ggml_tensor_overhead() * 16384 + ggml_graph_overhead(); - struct ggml_init_params gp = {bs, nullptr, true}; - auto* ctx0 = ggml_init(gp); - if (!ctx0) return false; - - auto* mask_in = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, INTERPOL, INTERPOL, 1, 1); - ggml_set_name(mask_in, "mem_mask"); - ggml_set_input(mask_in); - - // Learned mask downsampler: 4 stages conv+LN+GELU + final 1×1 - // Conv output is in ggml "conv" layout. permute(1,2,0,3) converts to "internal" layout - // where ne[0]=channel for LN2d. permute(2,0,1,3) converts back to "conv" layout. - auto* ds = mask_in; - for (int s = 0; s < 4; ++s) { - int out_ch = (int)model.mem_enc.ds_conv_w[s]->ne[3]; - ds = ggml_conv_2d(ctx0, model.mem_enc.ds_conv_w[s], ds, 2, 2, 1, 1, 1, 1); - ds = ggml_add(ctx0, ds, ggml_reshape_4d(ctx0, model.mem_enc.ds_conv_b[s], 1, 1, out_ch, 1)); - ds = ggml_cont(ctx0, ggml_permute(ctx0, ds, 1, 2, 0, 3)); // conv→internal (ne[0]=channel) - ds = sam3_layer_norm_2d(ctx0, ds, model.mem_enc.ds_norm_w[s], model.mem_enc.ds_norm_b[s]); - ds = ggml_gelu(ctx0, ds); - ds = ggml_cont(ctx0, ggml_permute(ctx0, ds, 2, 0, 1, 3)); // internal→conv - } - ds = ggml_conv_2d(ctx0, model.mem_enc.ds_conv_w[4], ds, 1, 1, 0, 0, 1, 1); - ds = ggml_add(ctx0, ds, ggml_reshape_4d(ctx0, model.mem_enc.ds_conv_b[4], 1, 1, D, 1)); - ds = ggml_cont(ctx0, ggml_permute(ctx0, ds, 1, 2, 0, 3)); // conv→internal [D, ...] - - // Pixel projection — use fresh input tensor to avoid pulling in the - // entire ViT+neck recomputation from state.neck_trk[2]'s dependency tree - auto* pix_in_raw = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, D, H, H, 1); - ggml_set_name(pix_in_raw, "mem_pix_feat"); - ggml_set_input(pix_in_raw); - auto* pix_in = ggml_cont(ctx0, ggml_permute(ctx0, pix_in_raw, 2, 0, 1, 3)); - auto* pix = ggml_conv_2d(ctx0, model.mem_enc.pix_proj_w, pix_in, 1, 1, 0, 0, 1, 1); - pix = ggml_add(ctx0, pix, ggml_reshape_4d(ctx0, model.mem_enc.pix_proj_b, 1, 1, D, 1)); - pix = ggml_cont(ctx0, ggml_permute(ctx0, pix, 1, 2, 0, 3)); - - // Fusion: ADD (not multiply) - auto* fused = ggml_add(ctx0, pix, ds); - for (int i = 0; i < 2; ++i) - fused = sam3_cxblock_forward(ctx0, fused, - model.mem_enc.fuser_dw_w[i], model.mem_enc.fuser_dw_b[i], - model.mem_enc.fuser_norm_w[i], model.mem_enc.fuser_norm_b[i], - model.mem_enc.fuser_fc1_w[i], model.mem_enc.fuser_fc1_b[i], - model.mem_enc.fuser_fc2_w[i], model.mem_enc.fuser_fc2_b[i], - model.mem_enc.fuser_gamma[i]); - auto* fused_out = ggml_cont(ctx0, ggml_permute(ctx0, fused, 2, 0, 1, 3)); - auto* mo = ggml_conv_2d(ctx0, model.mem_enc.out_proj_w, fused_out, 1, 1, 0, 0, 1, 1); - mo = ggml_add(ctx0, mo, ggml_reshape_4d(ctx0, model.mem_enc.out_proj_b, 1, 1, MD, 1)); - mo = ggml_cont(ctx0, ggml_permute(ctx0, mo, 1, 2, 0, 3)); - ggml_set_name(mo, "mem_out"); - ggml_set_output(mo); - - auto* g = ggml_new_graph_custom(ctx0, 16384, false); - ggml_build_forward_expand(g, mo); - auto* ga = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); - if (!ggml_gallocr_reserve(ga, g) || !ggml_gallocr_alloc_graph(ga, g)) { - ggml_gallocr_free(ga); - ggml_free(ctx0); - return false; - } - ggml_backend_tensor_set(mask_in, m_interp.data(), 0, m_interp.size() * sizeof(float)); - // Copy pixel features from state tensor to fresh input - { - std::vector pix_data(D * H * H); - ggml_backend_tensor_get(state.neck_trk[2], pix_data.data(), 0, D * H * H * sizeof(float)); - ggml_backend_tensor_set(pix_in_raw, pix_data.data(), 0, D * H * H * sizeof(float)); - } - if (!sam3_graph_compute(model.backend, g, 4)) { - ggml_gallocr_free(ga); - ggml_free(ctx0); - return false; - } - - std::vector md(MD * H * H); - ggml_backend_tensor_get(mo, md.data(), 0, md.size() * sizeof(float)); - - // Apply no_obj_embed_spatial if occluded (SAM2.1 only — EdgeTAM does not have this) - if (obj_score <= 0.0f && model.no_obj_embed_spatial) { - std::vector no_obj_emb(MD); - auto* noe = model.no_obj_embed_spatial; - if (noe->type == GGML_TYPE_F16) { - std::vector tmp(MD); - ggml_backend_tensor_get(noe, tmp.data(), 0, MD * sizeof(ggml_fp16_t)); - ggml_fp16_to_fp32_row(tmp.data(), no_obj_emb.data(), MD); - } else { - ggml_backend_tensor_get(noe, no_obj_emb.data(), 0, MD * sizeof(float)); - } - for (int i = 0; i < MD * H * H; ++i) - md[i] += no_obj_emb[i % MD]; - } - - if (!tracker.ctx) { - struct ggml_init_params tp = {ggml_tensor_overhead() * 4096, nullptr, true}; - tracker.ctx = ggml_init(tp); - } - - // ── EdgeTAM perceiver: compress memory features before storage ─────── - if (hp.has_perceiver) { - // Compute sinusoidal PE for memory features (needed as perceiver input) - sam3_ensure_tracker_pe_caches(tracker, hp, H); - const auto& mem_pos = tracker.cached_sinpe_64; - - std::vector perc_latents, perc_pos; - if (!edgetam_perceiver_forward(model, md, mem_pos, H, H, - perc_latents, perc_pos)) { - fprintf(stderr, "%s: perceiver forward failed\n", __func__); - ggml_gallocr_free(ga); - ggml_free(ctx0); - return false; - } - - // Store perceiver output: [MD=64, N_perc] (N_1d + N_2d latents) - const int N_perc = hp.perceiver_n_latents_1d + hp.perceiver_n_latents_2d; - auto* st = ggml_new_tensor_2d(tracker.ctx, GGML_TYPE_F32, MD, N_perc); - auto* sb = ggml_backend_alloc_buffer(model.backend, MD * N_perc * sizeof(float)); - struct ggml_tallocr ta_perc = ggml_tallocr_new(sb); - ggml_tallocr_alloc(&ta_perc, st); - tracker.owned_buffers.push_back(sb); - ggml_backend_tensor_set(st, perc_latents.data(), 0, MD * N_perc * sizeof(float)); - - // Store perceiver PE: [MD=64, 512] (first 256 zeros, last 256 sinusoidal) - auto* spe = ggml_new_tensor_2d(tracker.ctx, GGML_TYPE_F32, MD, N_perc); - auto* speb = ggml_backend_alloc_buffer(model.backend, MD * N_perc * sizeof(float)); - struct ggml_tallocr ta_perc2 = ggml_tallocr_new(speb); - ggml_tallocr_alloc(&ta_perc2, spe); - tracker.owned_buffers.push_back(speb); - ggml_backend_tensor_set(spe, perc_pos.data(), 0, MD * N_perc * sizeof(float)); - - sam3_memory_slot slot; - slot.spatial_feats = st; - slot.spatial_pe = spe; - slot.frame_index = frame_idx; - slot.is_cond_frame = is_cond; - auto& bk = tracker.mem_banks[inst_id]; - bk.push_back(slot); - while ((int)bk.size() > hp.num_maskmem) { - bool removed = false; - for (auto it = bk.begin(); it != bk.end(); ++it) - if (!it->is_cond_frame) { - bk.erase(it); - removed = true; - break; - } - if (!removed) bk.erase(bk.begin() + 1); - } - ggml_gallocr_free(ga); - ggml_free(ctx0); - return true; - } - - // ── Standard path (SAM2/SAM3): store raw [MD, H, H] features ──────── - // Store spatial features - auto* st = ggml_new_tensor_4d(tracker.ctx, GGML_TYPE_F32, MD, H, H, 1); - auto* sb = ggml_backend_alloc_buffer(model.backend, MD * H * H * sizeof(float)); - struct ggml_tallocr ta = ggml_tallocr_new(sb); - ggml_tallocr_alloc(&ta, st); - tracker.owned_buffers.push_back(sb); - ggml_backend_tensor_set(st, md.data(), 0, md.size() * sizeof(float)); - - // Compute and store sinusoidal spatial PE - sam3_ensure_tracker_pe_caches(tracker, hp, H); - const auto& pe_data = tracker.cached_sinpe_64; - auto* spe = ggml_new_tensor_4d(tracker.ctx, GGML_TYPE_F32, MD, H, H, 1); - auto* speb = ggml_backend_alloc_buffer(model.backend, MD * H * H * sizeof(float)); - struct ggml_tallocr ta2 = ggml_tallocr_new(speb); - ggml_tallocr_alloc(&ta2, spe); - tracker.owned_buffers.push_back(speb); - ggml_backend_tensor_set(spe, pe_data.data(), 0, pe_data.size() * sizeof(float)); - - sam3_memory_slot slot; - slot.spatial_feats = st; - slot.spatial_pe = spe; - slot.frame_index = frame_idx; - slot.is_cond_frame = is_cond; - auto& bk = tracker.mem_banks[inst_id]; - bk.push_back(slot); - while ((int)bk.size() > hp.num_maskmem) { - bool removed = false; - for (auto it = bk.begin(); it != bk.end(); ++it) - if (!it->is_cond_frame) { - bk.erase(it); - removed = true; - break; - } - if (!removed) bk.erase(bk.begin() + 1); - } - ggml_gallocr_free(ga); - ggml_free(ctx0); - return true; -} - -static void sam3_store_obj_ptr( - sam3_tracker& tracker, const sam3_model& model, - int inst_id, const float* pd, int frame_idx) { - const int D = model.hparams.neck_dim; - if (!tracker.ctx) { - struct ggml_init_params tp = {ggml_tensor_overhead() * 4096, nullptr, true}; - tracker.ctx = ggml_init(tp); - } - auto* pt = ggml_new_tensor_2d(tracker.ctx, GGML_TYPE_F32, D, 1); - auto* pb = ggml_backend_alloc_buffer(model.backend, D * sizeof(float)); - struct ggml_tallocr ta = ggml_tallocr_new(pb); - ggml_tallocr_alloc(&ta, pt); - tracker.owned_buffers.push_back(pb); - ggml_backend_tensor_set(pt, pd, 0, D * sizeof(float)); - auto& bk = tracker.ptr_banks[inst_id]; - bk.push_back({frame_idx, pt}); - while ((int)bk.size() > model.hparams.max_obj_ptrs) bk.erase(bk.begin()); -} - -sam3_tracker_ptr sam3_create_tracker(const sam3_model& model, - const sam3_video_params& params) { - if (model.hparams.is_sam2()) { - fprintf(stderr, "%s: ERROR: text-prompted tracker not available for SAM2 " - "(use sam3_create_visual_tracker instead)\n", __func__); - return nullptr; - } - sam3_tracker_ptr tracker(new sam3_tracker()); - tracker->params = params; - fprintf(stderr, "%s: tracker created (hotstart=%d, max_keep_alive=%d)\n", - __func__, params.hotstart_delay, params.max_keep_alive); - return tracker; -} - -sam3_result sam3_track_frame(sam3_tracker& tracker, sam3_state& state, - const sam3_model& model, const sam3_image& frame) { - if (model.hparams.visual_only) { - fprintf(stderr, "%s: ERROR: track_frame not available on visual-only model " - "(use sam3_propagate_frame instead)\n", __func__); - return sam3_result{}; - } - - sam3_result result; - const int D = model.hparams.neck_dim; - if (!sam3_encode_image(state, model, frame)) return result; - int fi = tracker.frame_index; - fprintf(stderr, "%s: frame %d (%zu active + %zu pending)\n", - __func__, fi, tracker.masklets.size(), tracker.pending.size()); - - std::map pm; - std::map po; - for (auto& ml : tracker.masklets) { - int id = ml.instance_id; - auto im = tracker.mem_banks.find(id); - if (im == tracker.mem_banks.end() || im->second.empty()) continue; - po[id] = sam3_propagate_single(tracker, state, model, ml, im->second, tracker.ptr_banks[id]); - if (po[id].mask_logits.empty()) continue; - auto rs = sam3_bilinear_interpolate(po[id].mask_logits.data(), - po[id].mask_w, po[id].mask_h, state.orig_width, state.orig_height); - pm[id].width = state.orig_width; - pm[id].height = state.orig_height; - pm[id].data.resize(state.orig_width * state.orig_height); - int fg = 0; - for (int p = 0; p < (int)rs.size(); ++p) { - bool f = rs[p] > 0.0f; - pm[id].data[p] = f ? 255 : 0; - if (f) fg++; - } - ml.last_score = po[id].iou_scores[0]; - ml.last_seen = fi; - float cov = (float)fg / (state.orig_width * state.orig_height); - ml.mds_sum += (cov > 0.001f && po[id].obj_score > 0.0f) ? 1 : -1; - } - for (auto& ml : tracker.pending) { - int id = ml.instance_id; - auto im = tracker.mem_banks.find(id); - if (im == tracker.mem_banks.end() || im->second.empty()) continue; - auto p2 = sam3_propagate_single(tracker, state, model, ml, im->second, tracker.ptr_banks[id]); - if (!p2.mask_logits.empty()) { - ml.last_score = p2.iou_scores[0]; - ml.last_seen = fi; - auto r2 = sam3_bilinear_interpolate(p2.mask_logits.data(), - p2.mask_w, p2.mask_h, state.orig_width, state.orig_height); - int fg2 = 0; - for (auto v : r2) - if (v > 0.0f) fg2++; - float c2 = (float)fg2 / (state.orig_width * state.orig_height); - ml.mds_sum += (c2 > 0.001f && p2.obj_score > 0.0f) ? 1 : -1; - sam3_encode_memory(tracker, state, model, id, - p2.mask_logits.data(), p2.mask_h, p2.mask_w, fi, false, p2.obj_score); - std::vector op(D); - sam3_extract_obj_ptr_cpu(model, p2.sam_token.data(), p2.obj_score, op.data()); - sam3_store_obj_ptr(tracker, model, id, op.data(), fi); - } - } - sam3_result nd; - if (!tracker.params.text_prompt.empty()) { - sam3_pcs_params pcs; - pcs.text_prompt = tracker.params.text_prompt; - pcs.score_threshold = tracker.params.score_threshold; - pcs.nms_threshold = tracker.params.nms_threshold; - nd = sam3_segment_pcs(state, model, pcs); - } - // Match new PCS detections against ALL tracked objects (both active AND pending) - // so that the same object isn't detected as a new instance every frame. - std::vector all_tracked; - std::vector all_tracked_masks; - for (auto& ml : tracker.masklets) { - all_tracked.push_back(ml); - auto it = pm.find(ml.instance_id); - all_tracked_masks.push_back(it != pm.end() ? it->second : sam3_mask{}); - } - for (auto& ml : tracker.pending) { - all_tracked.push_back(ml); - auto it = pm.find(ml.instance_id); - all_tracked_masks.push_back(it != pm.end() ? it->second : sam3_mask{}); - } - auto mat = sam3_match_detections(all_tracked, nd.detections, - all_tracked_masks, tracker.params.assoc_iou_threshold); - std::vector dmat(nd.detections.size(), false); - for (auto& m : mat) { - dmat[m.second] = true; - // Update last_seen on the matched masklet (in active or pending) - int mid = all_tracked[m.first].instance_id; - for (auto& ml : tracker.masklets) - if (ml.instance_id == mid) { - ml.last_seen = fi; - break; - } - for (auto& ml : tracker.pending) - if (ml.instance_id == mid) { - ml.last_seen = fi; - ml.mds_sum++; - break; - } - } - // Map: pending instance_id → PCS detection index (for display) - std::map pending_det_idx; - - for (int j = 0; j < (int)nd.detections.size(); ++j) { - if (dmat[j]) continue; - sam3_masklet ml; - ml.instance_id = tracker.next_inst_id++; - ml.first_frame = fi; - ml.last_seen = fi; - ml.last_score = nd.detections[j].score; - ml.mds_sum = 1; - pending_det_idx[ml.instance_id] = j; - - // Encode the detection's mask into memory so subsequent frames can - // propagate from it. Without this, pending masklets have no memory - // and propagation returns empty on all future frames. - const auto& det = nd.detections[j]; - if (!det.mask.data.empty()) { - // Build mask logits from the PCS detection's binary mask. - const int mh = sam3_eff_feat_size(state, model.hparams) * 4, mw = mh; - // Resize to mh×mw logit space: inside mask → +5.0, outside → -5.0. - std::vector det_logits(mh * mw); - for (int r = 0; r < mh; ++r) { - for (int c = 0; c < mw; ++c) { - int sx = c * det.mask.width / mw; - int sy = r * det.mask.height / mh; - sx = std::min(sx, det.mask.width - 1); - sy = std::min(sy, det.mask.height - 1); - det_logits[r * mw + c] = (det.mask.data[sy * det.mask.width + sx] > 127) - ? 5.0f - : -5.0f; - } - } - sam3_encode_memory(tracker, state, model, ml.instance_id, - det_logits.data(), mh, mw, fi, true, 1.0f); - - // Store a dummy object pointer (from no_obj_ptr since we don't - // have a SAM token from PCS — PCS uses a different decoder) - std::vector no_ptr(D); - auto* nop = model.tensors.at("no_obj_ptr"); - if (nop->type == GGML_TYPE_F16) { - std::vector tmp(D); - ggml_backend_tensor_get(nop, tmp.data(), 0, D * sizeof(ggml_fp16_t)); - ggml_fp16_to_fp32_row(tmp.data(), no_ptr.data(), D); - } else { - ggml_backend_tensor_get(nop, no_ptr.data(), 0, D * sizeof(float)); - } - sam3_store_obj_ptr(tracker, model, ml.instance_id, no_ptr.data(), fi); - } - - tracker.pending.push_back(std::move(ml)); - } - for (auto& ml : tracker.masklets) { - int id = ml.instance_id; - auto it = po.find(id); - if (it == po.end() || it->second.mask_logits.empty()) continue; - sam3_encode_memory(tracker, state, model, id, - it->second.mask_logits.data(), it->second.mask_h, it->second.mask_w, fi, false, it->second.obj_score); - std::vector op(D); - sam3_extract_obj_ptr_cpu(model, it->second.sam_token.data(), it->second.obj_score, op.data()); - sam3_store_obj_ptr(tracker, model, id, op.data(), fi); - } - sam3_update_tracker(tracker, fi); - - // Helper: build detection from a mask and add to result - auto add_mask_to_result = [&](int inst_id, float score, const sam3_mask& mask) { - if (mask.data.empty()) return; - sam3_detection det; - det.instance_id = inst_id; - det.score = score; - det.mask = mask; - det.mask.instance_id = inst_id; - det.mask.iou_score = score; - float x0 = 1e9f, y0 = 1e9f, x1 = -1e9f, y1 = -1e9f; - for (int p = 0; p < (int)det.mask.data.size(); ++p) - if (det.mask.data[p] > 127) { - int x = p % det.mask.width, y = p / det.mask.width; - x0 = std::min(x0, (float)x); - y0 = std::min(y0, (float)y); - x1 = std::max(x1, (float)x); - y1 = std::max(y1, (float)y); - } - if (x0 <= x1) det.box = {x0, y0, x1, y1}; - result.detections.push_back(std::move(det)); - }; - - // Include confirmed (active) masklet propagation results - for (auto& ml : tracker.masklets) { - auto it = pm.find(ml.instance_id); - if (it != pm.end()) add_mask_to_result(ml.instance_id, ml.last_score, it->second); - } - - // Include pending masklet results: on the detection frame, use the PCS - // detection mask; on subsequent frames, use the propagated mask. - for (auto& ml : tracker.pending) { - auto it = pm.find(ml.instance_id); - if (it != pm.end() && !it->second.data.empty()) { - // Propagated mask available (frame > detection frame) - add_mask_to_result(ml.instance_id, ml.last_score, it->second); - } else { - // First detection frame: use the PCS detection's mask+box directly - auto di = pending_det_idx.find(ml.instance_id); - if (di != pending_det_idx.end() && di->second < (int)nd.detections.size()) { - const auto& det = nd.detections[di->second]; - if (!det.mask.data.empty()) { - add_mask_to_result(ml.instance_id, ml.last_score, det.mask); - } - } - } - } - sam3_resolve_overlaps(result.detections); - for (auto& d : result.detections) { - if (d.mask.data.empty()) continue; - sam3_fill_holes(d.mask.data.data(), d.mask.width, d.mask.height, tracker.params.fill_hole_area); - sam3_remove_sprinkles(d.mask.data.data(), d.mask.width, d.mask.height, tracker.params.fill_hole_area); - } - tracker.frame_index++; - SAM3_LOG(2, "%s: frame %d done — %zu tracked\n", __func__, fi, result.detections.size()); - return result; -} - -bool sam3_refine_instance(sam3_tracker& tracker, sam3_state& state, - const sam3_model& model, int instance_id, - const std::vector& pos_points, - const std::vector& neg_points) { - const int D = model.hparams.neck_dim; - sam3_masklet* tgt = nullptr; - for (auto& ml : tracker.masklets) - if (ml.instance_id == instance_id) { - tgt = &ml; - break; - } - if (!tgt) - for (auto& ml : tracker.pending) - if (ml.instance_id == instance_id) { - tgt = &ml; - break; - } - if (!tgt) { - fprintf(stderr, "%s: instance %d not found\n", __func__, instance_id); - return false; - } - sam3_pvs_params pvs; - pvs.pos_points = pos_points; - pvs.neg_points = neg_points; - pvs.multimask = false; - auto r = sam3_segment_pvs(state, model, pvs); - if (r.detections.empty()) return false; - // tracker.frame_index points to the *next* frame; the refinement applies - // to the frame that was last tracked / encoded. - int fi = std::max(0, tracker.frame_index - 1); - const auto& rdet = r.detections[0]; - tgt->last_score = rdet.score; - tgt->last_seen = fi; - std::vector op(D); - if (!rdet.sam_token.empty()) { - sam3_extract_obj_ptr_cpu(model, rdet.sam_token.data(), rdet.mask.obj_score, op.data()); + state.pe_gauss_cache.resize(pe_nel); + if (pe.pe_gaussian->type == GGML_TYPE_F16) { + std::vector tmp(pe_nel); + ggml_backend_tensor_get(pe.pe_gaussian, tmp.data(), 0, pe_nel * sizeof(ggml_fp16_t)); + ggml_fp16_to_fp32_row(tmp.data(), state.pe_gauss_cache.data(), pe_nel); } else { - std::fill(op.begin(), op.end(), 0.0f); - } - sam3_store_obj_ptr(tracker, model, instance_id, op.data(), fi); - SAM3_LOG(2, "%s: refined instance %d\n", __func__, instance_id); - return true; -} - -int sam3_tracker_add_instance(sam3_tracker& tracker, sam3_state& state, - const sam3_model& model, - const sam3_pvs_params& pvs_params) { - const int D = model.hparams.neck_dim; - const int mask_hw = sam3_eff_feat_size(state, model.hparams) * 4; + ggml_backend_tensor_get(pe.pe_gaussian, state.pe_gauss_cache.data(), 0, pe_nel * sizeof(float)); + } - // Run PVS to get the segmentation mask - auto r = sam3_segment_pvs(state, model, pvs_params); - if (r.detections.empty()) { - fprintf(stderr, "%s: PVS returned no masks\n", __func__); - return -1; + for (int i = 0; i < 4; ++i) { + if (pe.point_embed[i]->type == GGML_TYPE_F16) { + std::vector tmp(D); + ggml_backend_tensor_get(pe.point_embed[i], tmp.data(), 0, D * sizeof(ggml_fp16_t)); + ggml_fp16_to_fp32_row(tmp.data(), state.point_emb_cache[i], D); + } else { + ggml_backend_tensor_get(pe.point_embed[i], state.point_emb_cache[i], 0, D * sizeof(float)); + } } - const auto& det = r.detections[0]; - if (det.mask.data.empty()) { - fprintf(stderr, "%s: PVS mask is empty\n", __func__); - return -1; + if (pe.not_a_point_embed->type == GGML_TYPE_F16) { + std::vector tmp(D); + ggml_backend_tensor_get(pe.not_a_point_embed, tmp.data(), 0, D * sizeof(ggml_fp16_t)); + ggml_fp16_to_fp32_row(tmp.data(), state.not_a_point_cache, D); + } else { + ggml_backend_tensor_get(pe.not_a_point_embed, state.not_a_point_cache, 0, D * sizeof(float)); } - int inst_id = tracker.next_inst_id++; - // tracker.frame_index points to the *next* frame to process; - // the instance is being added on the frame that was just tracked. - int fi = tracker.frame_index - 1; - if (fi < 0) fi = 0; + if (pe.no_mask_embed->type == GGML_TYPE_F16) { + std::vector tmp(D); + ggml_backend_tensor_get(pe.no_mask_embed, tmp.data(), 0, D * sizeof(ggml_fp16_t)); + ggml_fp16_to_fp32_row(tmp.data(), state.no_mask_emb_cache, D); + } else { + ggml_backend_tensor_get(pe.no_mask_embed, state.no_mask_emb_cache, 0, D * sizeof(float)); + } - // Create synthetic 288x288 logits from the binary mask. - // sam3_encode_memory applies sigmoid then scale/bias, so +6/-6 gives - // sigmoid values ~0.9975/~0.0025 — practically identical to real logits. - std::vector synth_logits(mask_hw * mask_hw); - { - int mw = det.mask.width, mh = det.mask.height; - for (int y = 0; y < mask_hw; ++y) { - int sy = y * mh / mask_hw; - for (int x = 0; x < mask_hw; ++x) { - int sx = x * mw / mask_hw; - synth_logits[y * mask_hw + x] = - (det.mask.data[sy * mw + sx] > 127) ? 6.0f : -6.0f; - } + state.dense_pe_cache.resize(D * H * H); + for (int row = 0; row < H; ++row) { + for (int col = 0; col < H; ++col) { + float x_norm = ((float)col + 0.5f) / (float)H; + float y_norm = ((float)row + 0.5f) / (float)H; + float pe_vec[256]; + sam3_pe_encode_coord(pe_vec, x_norm, y_norm, + state.pe_gauss_cache.data(), num_pos_feats); + for (int d = 0; d < D; ++d) + state.dense_pe_cache[d + col * D + row * D * H] = pe_vec[d]; } } - // Encode into memory bank - float obj_score = det.mask.obj_score; - if (!sam3_encode_memory(tracker, state, model, inst_id, - synth_logits.data(), mask_hw, mask_hw, - fi, true, obj_score)) { - fprintf(stderr, "%s: failed to encode memory for instance %d\n", __func__, inst_id); - return -1; + state.dense_nomask_cache.resize(D * H * H); + for (int row = 0; row < H; ++row) { + for (int col = 0; col < H; ++col) { + for (int d = 0; d < D; ++d) + state.dense_nomask_cache[d + col * D + row * D * H] = state.no_mask_emb_cache[d]; + } } - // Extract object pointer from the SAM decoder token - std::vector op(D); - if (!det.sam_token.empty()) { - sam3_extract_obj_ptr_cpu(model, det.sam_token.data(), obj_score, op.data()); - } else { - std::fill(op.begin(), op.end(), 0.0f); - } - sam3_store_obj_ptr(tracker, model, inst_id, op.data(), fi); + state.pe_cache_valid = true; + fprintf(stderr, "%s: PE cache populated (%d embeddings, %.1f KB dense grids)\n", + __func__, pe_nel, 2.0f * D * H * H * sizeof(float) / 1024.0f); +} - // Create confirmed masklet - sam3_masklet ml; - ml.instance_id = inst_id; - ml.first_frame = fi; - ml.last_seen = fi; - ml.last_score = det.score; - ml.confirmed = true; - ml.mds_sum = 1; - tracker.masklets.push_back(std::move(ml)); +// Match the official image predictor prompt ordering: +// box corners (if any) → user points → trailing padding point. +// Even when a box is present, SAM keeps the final padding token because boxes +// are merged into the point stream before calling PromptEncoder(points=..., boxes=None). +static void sam3_collect_pvs_prompt_tokens(const sam3_pvs_params& params, + std::vector& all_coords, + std::vector& all_labels) { + all_coords.clear(); + all_labels.clear(); - SAM3_LOG(2, "%s: added instance #%d (score=%.3f)\n", __func__, inst_id, det.score); - return inst_id; -} + if (params.use_box) { + all_coords.push_back(params.box.x0); + all_coords.push_back(params.box.y0); + all_labels.push_back(2); -int sam3_tracker_frame_index(const sam3_tracker& tracker) { return tracker.frame_index; } + all_coords.push_back(params.box.x1); + all_coords.push_back(params.box.y1); + all_labels.push_back(3); + } -void sam3_tracker_reset(sam3_tracker& tracker) { - tracker.frame_index = 0; - tracker.next_inst_id = 1; - tracker.masklets.clear(); - tracker.pending.clear(); - tracker.mem_banks.clear(); - tracker.ptr_banks.clear(); - for (auto* b : tracker.owned_buffers) - if (b) ggml_backend_buffer_free(b); - tracker.owned_buffers.clear(); - if (tracker.ctx) { - ggml_free(tracker.ctx); - tracker.ctx = nullptr; + for (const auto& pt : params.pos_points) { + all_coords.push_back(pt.x); + all_coords.push_back(pt.y); + all_labels.push_back(1); } - if (tracker.buffer) { - ggml_backend_buffer_free(tracker.buffer); - tracker.buffer = nullptr; + for (const auto& pt : params.neg_points) { + all_coords.push_back(pt.x); + all_coords.push_back(pt.y); + all_labels.push_back(0); } -} - -/***************************************************************************** -** Visual-only video tracking -*****************************************************************************/ -sam3_tracker_ptr sam3_create_visual_tracker( - const sam3_model& model, - const sam3_visual_track_params& params) { - sam3_video_params vp; - vp.text_prompt = ""; // no PCS detection - vp.assoc_iou_threshold = params.assoc_iou_threshold; - vp.max_keep_alive = params.max_keep_alive; - vp.recondition_every = params.recondition_every; - vp.fill_hole_area = params.fill_hole_area; - sam3_tracker_ptr tracker(new sam3_tracker()); - tracker->params = vp; - fprintf(stderr, "%s: visual-only tracker created (max_keep_alive=%d)\n", - __func__, params.max_keep_alive); - return tracker; + all_coords.push_back(0.0f); + all_coords.push_back(0.0f); + all_labels.push_back(-1); } -sam3_result sam3_propagate_frame( - sam3_tracker& tracker, sam3_state& state, - const sam3_model& model, const sam3_image& frame) { - sam3_result result; - const int D = model.hparams.neck_dim; - if (!sam3_encode_image(state, model, frame)) return result; - int fi = tracker.frame_index; - fprintf(stderr, "%s: frame %d (%zu active + %zu pending)\n", - __func__, fi, tracker.masklets.size(), tracker.pending.size()); - - // ── Propagate active masklets ──────────────────────────────────────── - std::map pm; - std::map po; - for (auto& ml : tracker.masklets) { - int id = ml.instance_id; - auto im = tracker.mem_banks.find(id); - if (im == tracker.mem_banks.end() || im->second.empty()) continue; - po[id] = sam3_propagate_single(tracker, state, model, ml, im->second, tracker.ptr_banks[id]); - if (po[id].mask_logits.empty()) continue; - auto rs = sam3_bilinear_interpolate(po[id].mask_logits.data(), - po[id].mask_w, po[id].mask_h, - state.orig_width, state.orig_height); - pm[id].width = state.orig_width; - pm[id].height = state.orig_height; - pm[id].data.resize(state.orig_width * state.orig_height); - int fg = 0; - for (int p = 0; p < (int)rs.size(); ++p) { - bool f = rs[p] > 0.0f; - pm[id].data[p] = f ? 255 : 0; - if (f) fg++; - } - ml.last_score = po[id].iou_scores[0]; - ml.last_seen = fi; - float cov = (float)fg / (state.orig_width * state.orig_height); - ml.mds_sum += (cov > 0.001f && po[id].obj_score > 0.0f) ? 1 : -1; - } +// Build sparse and dense embeddings from point/box prompts +// sparse_out: [D, N_pts, 1] where N_pts = n_pos + n_neg + (use_box ? 2 : 0) + pad +// dense_out: [D, H, H, 1] (no-mask default or mask downsample) +struct sam3_pe_result { + struct ggml_tensor* sparse; // [D, N_pts, 1] + struct ggml_tensor* dense; // [D, H, H, 1] + struct ggml_tensor* image_pe; // [D, H, H, 1] — dense positional encoding grid + int n_tokens; +}; - // ── Propagate pending masklets ─────────────────────────────────────── - for (auto& ml : tracker.pending) { - int id = ml.instance_id; - auto im = tracker.mem_banks.find(id); - if (im == tracker.mem_banks.end() || im->second.empty()) continue; - auto p2 = sam3_propagate_single(tracker, state, model, ml, im->second, tracker.ptr_banks[id]); - if (!p2.mask_logits.empty()) { - ml.last_score = p2.iou_scores[0]; - ml.last_seen = fi; - auto r2 = sam3_bilinear_interpolate(p2.mask_logits.data(), - p2.mask_w, p2.mask_h, - state.orig_width, state.orig_height); - int fg2 = 0; - for (auto v : r2) - if (v > 0.0f) fg2++; - float c2 = (float)fg2 / (state.orig_width * state.orig_height); - ml.mds_sum += (c2 > 0.001f && p2.obj_score > 0.0f) ? 1 : -1; - pm[id].width = state.orig_width; - pm[id].height = state.orig_height; - pm[id].data.resize(state.orig_width * state.orig_height); - for (int p = 0; p < (int)r2.size(); ++p) - pm[id].data[p] = r2[p] > 0.0f ? 255 : 0; - sam3_encode_memory(tracker, state, model, id, - p2.mask_logits.data(), p2.mask_h, p2.mask_w, - fi, false, p2.obj_score); - std::vector op(D); - sam3_extract_obj_ptr_cpu(model, p2.sam_token.data(), p2.obj_score, op.data()); - sam3_store_obj_ptr(tracker, model, id, op.data(), fi); - } - } +static sam3_pe_result sam3_build_sam_pe( + struct ggml_context* ctx, + const sam3_pvs_params& params, + int embed_dim, int feat_size) { + const int D = embed_dim; // 256 + const int H = feat_size; // 72 - // ── Encode memory for active masklets ──────────────────────────────── - for (auto& ml : tracker.masklets) { - int id = ml.instance_id; - auto it = po.find(id); - if (it == po.end() || it->second.mask_logits.empty()) continue; - sam3_encode_memory(tracker, state, model, id, - it->second.mask_logits.data(), it->second.mask_h, - it->second.mask_w, fi, false, it->second.obj_score); - std::vector op(D); - sam3_extract_obj_ptr_cpu(model, it->second.sam_token.data(), - it->second.obj_score, op.data()); - sam3_store_obj_ptr(tracker, model, id, op.data(), fi); - } + int N_pts = (int)(params.pos_points.size() + params.neg_points.size()) + 1; // pad + if (params.use_box) N_pts += 2; - // ── Update tracker state (confirmation / eviction) ─────────────────── - sam3_update_tracker(tracker, fi); + auto* sparse = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, N_pts, 1); + ggml_set_name(sparse, "sam_pe_sparse"); + ggml_set_input(sparse); - // ── Build result ───────────────────────────────────────────────────── - auto add_mask_to_result = [&](int inst_id, float score, const sam3_mask& mask) { - if (mask.data.empty()) return; - sam3_detection det; - det.instance_id = inst_id; - det.score = score; - det.mask = mask; - det.mask.instance_id = inst_id; - det.mask.iou_score = score; - float x0 = 1e9f, y0 = 1e9f, x1 = -1e9f, y1 = -1e9f; - for (int p = 0; p < (int)det.mask.data.size(); ++p) - if (det.mask.data[p] > 127) { - int x = p % det.mask.width, y = p / det.mask.width; - x0 = std::min(x0, (float)x); - y0 = std::min(y0, (float)y); - x1 = std::max(x1, (float)x); - y1 = std::max(y1, (float)y); - } - if (x0 <= x1) det.box = {x0, y0, x1, y1}; - result.detections.push_back(std::move(det)); - }; + auto* image_pe = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, D, H, H, 1); + ggml_set_name(image_pe, "sam_pe_image_pe"); + ggml_set_input(image_pe); - for (auto& ml : tracker.masklets) { - auto it = pm.find(ml.instance_id); - if (it != pm.end()) add_mask_to_result(ml.instance_id, ml.last_score, it->second); - } - for (auto& ml : tracker.pending) { - auto it = pm.find(ml.instance_id); - if (it != pm.end()) add_mask_to_result(ml.instance_id, ml.last_score, it->second); - } + auto* dense = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, D, H, H, 1); + ggml_set_name(dense, "sam_pe_dense"); + ggml_set_input(dense); - sam3_resolve_overlaps(result.detections); - for (auto& d : result.detections) { - if (d.mask.data.empty()) continue; - sam3_fill_holes(d.mask.data.data(), d.mask.width, d.mask.height, - tracker.params.fill_hole_area); - sam3_remove_sprinkles(d.mask.data.data(), d.mask.width, d.mask.height, - tracker.params.fill_hole_area); - } - tracker.frame_index++; - SAM3_LOG(2, "%s: frame %d done — %zu tracked\n", - __func__, fi, result.detections.size()); + sam3_pe_result result; + result.sparse = sparse; + result.dense = dense; + result.image_pe = image_pe; + result.n_tokens = N_pts; return result; } /***************************************************************************** -** Utility — image I/O +** SAM mask decoder — graph building (Phase 6, Step 6.2) *****************************************************************************/ -sam3_image sam3_load_image(const std::string& path) { - sam3_image img; - int w, h, c; - uint8_t* data = stbi_load(path.c_str(), &w, &h, &c, 3); - if (!data) { - fprintf(stderr, "%s: failed to load '%s'\n", __func__, path.c_str()); - return img; +// TwoWayAttentionBlock forward +static void sam3_twoway_block_forward( + struct ggml_context* ctx, + struct ggml_tensor*& queries, // [D, N_q, B] — modified in place + struct ggml_tensor*& keys, // [D, N_kv, B] — modified in place + struct ggml_tensor* query_pe, // [D, N_q, B] + struct ggml_tensor* key_pe, // [D, N_kv, B] + const sam3_twoway_block& blk, + int n_heads, + bool skip_first_layer_pe) { + // 1. Self-attention on queries + if (skip_first_layer_pe) { + // Python: queries = self.self_attn(q=queries, k=queries, v=queries) + // No residual connection when skipping first layer PE + queries = sam3_sam_attention(ctx, queries, queries, queries, blk.self_attn, n_heads); + } else { + auto* q = ggml_add(ctx, queries, query_pe); + auto* attn_out = sam3_sam_attention(ctx, q, q, queries, blk.self_attn, n_heads); + queries = ggml_add(ctx, queries, attn_out); + } + queries = sam3_layer_norm(ctx, queries, blk.norm1_w, blk.norm1_b); + if (skip_first_layer_pe) { + ggml_set_name(queries, "dbg_twoway_skip_sa_norm"); + ggml_set_output(queries); } - img.width = w; - img.height = h; - img.channels = 3; - img.data.assign(data, data + w * h * 3); - stbi_image_free(data); - return img; -} - -bool sam3_save_mask(const sam3_mask& mask, const std::string& path) { - if (mask.data.empty()) return false; - return stbi_write_png(path.c_str(), mask.width, mask.height, 1, - mask.data.data(), mask.width) != 0; -} - -sam3_image sam3_decode_video_frame(const std::string& video_path, int frame_index) { - sam3_image img; - - // Use ffmpeg to extract a single frame as raw RGB (frame-accurate) - char cmd[1024]; - snprintf(cmd, sizeof(cmd), - "ffmpeg -nostdin -loglevel error -i \"%s\" " - "-vf \"select=eq(n\\,%d)\" -vsync vfr -frames:v 1 " - "-f rawvideo -pix_fmt rgb24 pipe:1 2>%s", - video_path.c_str(), frame_index, SAM3_NULL_DEV); - // First, get dimensions - char info_cmd[1024]; - snprintf(info_cmd, sizeof(info_cmd), - "ffprobe -v error -select_streams v:0 " - "-show_entries stream=width,height -of csv=p=0 \"%s\" 2>%s", - video_path.c_str(), SAM3_NULL_DEV); - FILE* fp = popen(info_cmd, SAM3_POPEN_READ); - if (!fp) return img; - int w = 0, h = 0; - if (fscanf(fp, "%d,%d", &w, &h) != 2) { - pclose(fp); - return img; + // 2. Cross-attention: tokens attending to image + { + auto* q = ggml_add(ctx, queries, query_pe); + auto* k = ggml_add(ctx, keys, key_pe); + auto* attn_out = sam3_sam_attention(ctx, q, k, keys, blk.ca_tok2img, n_heads); + queries = ggml_add(ctx, queries, attn_out); + queries = sam3_layer_norm(ctx, queries, blk.norm2_w, blk.norm2_b); + } + if (skip_first_layer_pe) { + ggml_set_name(queries, "dbg_twoway_skip_ca_tok2img"); + ggml_set_output(queries); } - pclose(fp); - img.width = w; - img.height = h; - img.channels = 3; - img.data.resize(w * h * 3); + // 3. MLP on queries (ReLU activation) + { + auto* mlp = ggml_mul_mat(ctx, blk.mlp_fc1_w, queries); + mlp = ggml_add(ctx, mlp, blk.mlp_fc1_b); + mlp = ggml_relu(ctx, mlp); + mlp = ggml_mul_mat(ctx, blk.mlp_fc2_w, mlp); + mlp = ggml_add(ctx, mlp, blk.mlp_fc2_b); + queries = ggml_add(ctx, queries, mlp); + queries = sam3_layer_norm(ctx, queries, blk.norm3_w, blk.norm3_b); + } + if (skip_first_layer_pe) { + ggml_set_name(queries, "dbg_twoway_skip_mlp"); + ggml_set_output(queries); + } - fp = popen(cmd, SAM3_POPEN_READ); - if (!fp) { - img.data.clear(); - return img; + // 4. Cross-attention: image attending to tokens + { + auto* q = ggml_add(ctx, queries, query_pe); + auto* k = ggml_add(ctx, keys, key_pe); + // Note: q and k are swapped — image (k) attends to tokens (q) + auto* attn_out = sam3_sam_attention(ctx, k, q, queries, blk.ca_img2tok, n_heads); + keys = ggml_add(ctx, keys, attn_out); + keys = sam3_layer_norm(ctx, keys, blk.norm4_w, blk.norm4_b); } - size_t nread = fread(img.data.data(), 1, img.data.size(), fp); - pclose(fp); - if (nread != img.data.size()) { - img.data.clear(); + if (skip_first_layer_pe) { + ggml_set_name(keys, "dbg_twoway_skip_img2tok"); + ggml_set_output(keys); } +} - return img; +// MLP forward: N layers with ReLU (except last), optional sigmoid on last +static struct ggml_tensor* sam3_mlp_forward( + struct ggml_context* ctx, + struct ggml_tensor* x, + struct ggml_tensor* const* weights, + struct ggml_tensor* const* biases, + int n_layers, + bool sigmoid_output = false) { + for (int i = 0; i < n_layers; ++i) { + x = ggml_mul_mat(ctx, weights[i], x); + x = ggml_add(ctx, x, biases[i]); + if (i < n_layers - 1) { + x = ggml_relu(ctx, x); + } + } + if (sigmoid_output) { + x = ggml_sigmoid(ctx, x); + } + return x; } -sam3_video_info sam3_get_video_info(const std::string& video_path) { - sam3_video_info info; +// Full SAM mask decoder graph +// Inputs: +// image_feats: [D, H, H, 1] — tracker neck features (scale 2 = 72×72) +// image_pe: [D, H, H, 1] — dense positional encoding +// sparse_emb: [D, N_pts, 1] — sparse prompt embeddings +// dense_emb: [D, H, H, 1] — dense prompt embeddings (no_mask default) +// feat_s0: [D, H0, H0, 1] — high-res features (scale 0 = 288×288) +// feat_s1: [D, H1, H1, 1] — mid-res features (scale 1 = 144×144) +// Outputs: sam3_dec_result with masks, iou_pred, obj_score, sam_token_out +struct sam3_dec_result { + struct ggml_tensor* masks; // [288*288, N_masks, 1] + struct ggml_tensor* iou_pred; // [N_masks, 1] + struct ggml_tensor* obj_score; // [1, 1] + struct ggml_tensor* sam_token; // [D, 1] — for object pointer + struct ggml_tensor* mask_tokens; // [D, N_masks, 1] — raw SAM mask tokens +}; - char cmd[1024]; - snprintf(cmd, sizeof(cmd), - "ffprobe -v error -select_streams v:0 " - "-show_entries stream=width,height,r_frame_rate,nb_frames " - "-of csv=p=0 \"%s\" 2>%s", - video_path.c_str(), SAM3_NULL_DEV); - FILE* fp = popen(cmd, SAM3_POPEN_READ); - if (!fp) return info; +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 - int w = 0, h = 0, num = 0, den = 1, nf = 0; - if (fscanf(fp, "%d,%d,%d/%d,%d", &w, &h, &num, &den, &nf) >= 4) { - info.width = w; - info.height = h; - info.fps = (den > 0) ? static_cast(num) / den : 0.0f; - info.n_frames = nf; + // ── 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); } - pclose(fp); - return info; -} - -/***************************************************************************** -** Tokenizer — standalone test API (does not require model weights) -*****************************************************************************/ + 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); -// Global tokenizer instance for the test API. -static sam3_bpe_tokenizer g_test_tokenizer; -static bool g_test_tokenizer_loaded = false; + const int N_tok = 6 + N_pts; -bool sam3_test_load_tokenizer(const std::string& model_path) { - std::ifstream fin(model_path, std::ios::binary); - if (!fin) return false; + 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); - // Read header - uint32_t magic; - int32_t version, ftype, n_tensors; - fin.read(reinterpret_cast(&magic), 4); - fin.read(reinterpret_cast(&version), 4); - fin.read(reinterpret_cast(&ftype), 4); - fin.read(reinterpret_cast(&n_tensors), 4); - if (magic != SAM3_MAGIC || version != SAM3_FILE_VERSION) return false; + auto* queries = tokens; + auto* keys = src; + auto* query_pe = tokens; // query PE = initial point embedding + auto* key_pe = pos_src; - // Skip hparams - sam3_hparams hp; - if (!sam3_load_hparams(fin, hp)) return false; - if (hp.visual_only) return false; + 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); + } - // Skip tensors - for (int t = 0; t < n_tensors; ++t) { - int32_t n_dims, name_len, dtype; - fin.read(reinterpret_cast(&n_dims), 4); - fin.read(reinterpret_cast(&name_len), 4); - fin.read(reinterpret_cast(&dtype), 4); - if (fin.fail()) return false; + // Final attention: tokens → image + { + auto* q = ggml_add(ctx, queries, query_pe); + auto* k = ggml_add(ctx, keys, key_pe); + auto* attn_out = sam3_sam_attention(ctx, q, k, keys, dec.final_attn, n_heads); + queries = ggml_add(ctx, queries, attn_out); + queries = sam3_layer_norm(ctx, queries, dec.final_norm_w, dec.final_norm_b); + ggml_set_name(queries, "sam_dec_final_queries"); + } - // Read shape to compute data size - int64_t n_el = 1; - std::vector shape(n_dims); - for (int i = 0; i < n_dims; ++i) { - int32_t d; - fin.read(reinterpret_cast(&d), 4); - shape[i] = d; - n_el *= d; - } + // Debug: mark transformer outputs + ggml_set_name(queries, "dbg_dec_queries_out"); + ggml_set_output(queries); + ggml_set_name(keys, "dbg_dec_keys_out"); + ggml_set_output(keys); - // Skip name - fin.seekg(name_len, std::ios::cur); + // ── Extract output tokens ──────────────────────────────────────────── + // With pred_obj_scores=True (6 tokens): obj(0), iou(1), masks(2..5) + // With pred_obj_scores=False (5 tokens): iou(0), masks(1..4) + const int s = has_obj_score ? 1 : 0; + auto* iou_token_out = ggml_view_3d(ctx, queries, D, 1, 1, + queries->nb[1], queries->nb[2], + s * queries->nb[1]); + iou_token_out = ggml_cont(ctx, iou_token_out); // [D, 1, 1] - // Skip padding to 32-byte alignment - size_t pos = fin.tellg(); - size_t pad = (32 - pos % 32) % 32; - if (pad > 0) fin.seekg(pad, std::ios::cur); + auto* mask_tokens_out = ggml_view_3d(ctx, queries, D, num_mask_tokens, 1, + queries->nb[1], queries->nb[2], + (s + 1) * queries->nb[1]); + mask_tokens_out = ggml_cont(ctx, mask_tokens_out); // [D, 4, 1] + ggml_set_name(mask_tokens_out, "sam_dec_mask_tokens"); - // Compute data size and skip - const ggml_type file_type = static_cast(dtype); - size_t bytes; - if (ggml_is_quantized(file_type)) { - const int64_t n_rows = n_el / shape[0]; - bytes = ggml_row_size(file_type, shape[0]) * n_rows; - } else { - const size_t elem_size = (file_type == GGML_TYPE_F16) ? 2 : 4; - bytes = n_el * elem_size; - } - fin.seekg(bytes, std::ios::cur); - if (fin.fail()) return false; + struct ggml_tensor* obj_in = nullptr; + if (has_obj_score) { + obj_in = ggml_view_3d(ctx, queries, D, 1, 1, + queries->nb[1], queries->nb[2], 0); + obj_in = ggml_cont(ctx, obj_in); // [D, 1, 1] } - // Read embedded tokenizer - if (!sam3_load_bpe_vocab_from_stream(fin, g_test_tokenizer)) return false; - g_test_tokenizer_loaded = true; - return true; -} - -std::vector sam3_test_tokenize(const std::string& text) { - if (!g_test_tokenizer_loaded) return {}; - return sam3_tokenize(g_test_tokenizer, text, 32); -} + // SAM output token = first mask token, used for object pointer + auto* sam_token = ggml_view_2d(ctx, queries, D, 1, + queries->nb[1], (s + 1) * queries->nb[1]); + sam_token = ggml_cont(ctx, sam_token); // [D, 1] + ggml_set_name(sam_token, "sam_dec_sam_token"); -static bool sam3_dump_tensor_to_path(struct ggml_tensor* t, - const std::string& tensor_name, - const std::string& output_path) { - if (!t) { - fprintf(stderr, "%s: tensor '%s' is null\n", __func__, tensor_name.c_str()); - return false; - } + // Upscale: [D, H*H, 1] → ConvTranspose → high-res masks + auto* src_img = ggml_reshape_4d(ctx, keys, D, H, H, 1); + src_img = ggml_cont(ctx, ggml_permute(ctx, src_img, 2, 0, 1, 3)); - int64_t numel = 1; - for (int i = 0; i < GGML_MAX_DIMS; ++i) { - if (t->ne[i] > 0) { - numel *= t->ne[i]; - } - } + auto* up1 = ggml_conv_transpose_2d_p0(ctx, sam3_conv_transpose_weight(ctx, dec.up1_w), src_img, 2); + up1 = ggml_add(ctx, up1, ggml_reshape_4d(ctx, dec.up1_b, 1, 1, ggml_nelements(dec.up1_b), 1)); - std::vector data(numel); - if (t->type != GGML_TYPE_F16 && t->type != GGML_TYPE_F32) { - fprintf(stderr, "%s: unsupported tensor type %d for '%s'\n", - __func__, (int)t->type, tensor_name.c_str()); - return false; - } + auto* fs1 = ggml_cont(ctx, ggml_permute(ctx, feat_s1, 2, 0, 1, 3)); + auto* hs1 = ggml_conv_2d_sk_p0(ctx, dec.conv_s1_w, fs1); // [W, H, 64, B] + hs1 = ggml_add(ctx, hs1, ggml_reshape_4d(ctx, dec.conv_s1_b, 1, 1, 64, 1)); + ggml_set_name(hs1, "sam_dec_feat_s1_proj"); - const int64_t ne0 = t->ne[0]; - const int64_t ne1 = t->ne[1]; - const int64_t ne2 = t->ne[2]; - const int64_t ne3 = t->ne[3]; + // Python: act1(ln1(dc1(src) + feat_s1)) — add before LayerNorm + up1 = ggml_add(ctx, up1, hs1); + up1 = ggml_cont(ctx, ggml_permute(ctx, up1, 1, 2, 0, 3)); + up1 = sam3_layer_norm_2d(ctx, up1, dec.up1_norm_w, dec.up1_norm_b); - if (ggml_is_contiguous(t)) { - if (t->type == GGML_TYPE_F16) { - std::vector f16_data(numel); - ggml_backend_tensor_get(t, f16_data.data(), 0, numel * sizeof(ggml_fp16_t)); - ggml_fp16_to_fp32_row(f16_data.data(), data.data(), numel); - } else { - ggml_backend_tensor_get(t, data.data(), 0, numel * sizeof(float)); - } - } else if (t->nb[0] == ggml_type_size(t->type)) { - // Serialize non-contiguous logical tensors in row-major ggml order. - if (t->type == GGML_TYPE_F16) { - std::vector row(ne0); - for (int64_t i3 = 0; i3 < ne3; ++i3) { - for (int64_t i2 = 0; i2 < ne2; ++i2) { - for (int64_t i1 = 0; i1 < ne1; ++i1) { - const size_t row_idx = ((size_t) i3 * ne2 * ne1 + (size_t) i2 * ne1 + (size_t) i1) * ne0; - const size_t offs = i3 * t->nb[3] + i2 * t->nb[2] + i1 * t->nb[1]; - ggml_backend_tensor_get(t, row.data(), offs, ne0 * sizeof(ggml_fp16_t)); - ggml_fp16_to_fp32_row(row.data(), data.data() + row_idx, ne0); - } - } - } - } else { - for (int64_t i3 = 0; i3 < ne3; ++i3) { - for (int64_t i2 = 0; i2 < ne2; ++i2) { - for (int64_t i1 = 0; i1 < ne1; ++i1) { - const size_t row_idx = ((size_t) i3 * ne2 * ne1 + (size_t) i2 * ne1 + (size_t) i1) * ne0; - const size_t offs = i3 * t->nb[3] + i2 * t->nb[2] + i1 * t->nb[1]; - ggml_backend_tensor_get(t, data.data() + row_idx, offs, ne0 * sizeof(float)); - } - } - } - } - } else { - fprintf(stderr, "%s: unsupported non-contiguous layout for '%s' (nb0=%llu)\n", - __func__, tensor_name.c_str(), (unsigned long long) t->nb[0]); - return false; - } + up1 = ggml_gelu_erf(ctx, up1); - { - std::ofstream f(output_path + ".bin", std::ios::binary); - if (!f) return false; - f.write(reinterpret_cast(data.data()), numel * sizeof(float)); - } + // Permute back to [W, H, C, B] for next deconv + up1 = ggml_cont(ctx, ggml_permute(ctx, up1, 2, 0, 1, 3)); // [144, 144, 64, 1] - { - std::ofstream f(output_path + ".shape"); - if (!f) return false; - int ndims = ggml_n_dims(t); - for (int i = 0; i < ndims; ++i) { - if (i > 0) f << ","; - f << t->ne[i]; - } - f << "\n"; - } + // dc2: ConvTranspose2d(64, 32, k=2, s=2) → [288, 288, 32, 1] + auto* up2 = ggml_conv_transpose_2d_p0(ctx, sam3_conv_transpose_weight(ctx, dec.up2_w), up1, 2); + up2 = ggml_add(ctx, up2, ggml_reshape_4d(ctx, dec.up2_b, 1, 1, ggml_nelements(dec.up2_b), 1)); - fprintf(stderr, "%s: dumped '%s' [", __func__, tensor_name.c_str()); - for (int i = 0; i < ggml_n_dims(t); ++i) { - if (i > 0) fprintf(stderr, ", "); - fprintf(stderr, "%lld", (long long)t->ne[i]); - } - fprintf(stderr, "] to %s\n", output_path.c_str()); + // conv_s0: 1x1 conv on feat_s0 (256→32). feat_s0 is [C, W, H, B] — permute for conv. + auto* fs0 = ggml_cont(ctx, ggml_permute(ctx, feat_s0, 2, 0, 1, 3)); // [W, H, C, B] + auto* hs0 = ggml_conv_2d_sk_p0(ctx, dec.conv_s0_w, fs0); // [W, H, 32, B] + hs0 = ggml_add(ctx, hs0, ggml_reshape_4d(ctx, dec.conv_s0_b, 1, 1, 32, 1)); + ggml_set_name(hs0, "sam_dec_feat_s0_proj"); - return true; -} + // Python: act2(dc2(upscaled_embedding) + feat_s0) — no LayerNorm here + up2 = ggml_add(ctx, up2, hs0); // both [W, H, 32, B] -static bool sam3_dump_raw_f32_to_path(const float* data, - const std::vector& shape, - const std::string& output_path) { - int64_t numel = 1; - for (int64_t d : shape) { - numel *= d; - } + // Permute to [C, W, H, B] for subsequent operations + up2 = ggml_cont(ctx, ggml_permute(ctx, up2, 1, 2, 0, 3)); // [32, 288, 288, 1] - { - std::ofstream f(output_path + ".bin", std::ios::binary); - if (!f) { - return false; - } - f.write(reinterpret_cast(data), numel * sizeof(float)); - } + // GELU activation (exact, matching Python nn.GELU) + up2 = ggml_gelu_erf(ctx, up2); - { - std::ofstream f(output_path + ".shape"); - if (!f) { - return false; - } - for (size_t i = 0; i < shape.size(); ++i) { - if (i > 0) { - f << ","; - } - f << shape[i]; - } - f << "\n"; - } + // up2: [32, 288, 288, 1] — this is our upscaled_embedding + ggml_set_name(up2, "sam_dec_upscaled"); - return true; -} + // ── Hypernetwork: predict masks ────────────────────────────────────── + // For each mask token i, pass through 3-layer MLP to get [32] vector + // Then dot product with upscaled_embedding [32, (H*4)^2] to get mask + const int H4 = H * 4; + auto* up_flat = ggml_reshape_3d(ctx, up2, 32, H4 * H4, 1); -static bool sam3_load_ref_f32_data(const std::string& path, - std::vector& data, - int expected_numel = -1) { - std::ifstream shape_f(path + ".shape"); - if (!shape_f) { - fprintf(stderr, "%s: missing %s.shape\n", __func__, path.c_str()); - return false; - } + // Process each mask token through its hypernetwork MLP + // mask_tokens_out: [D, 4, 1] + struct ggml_tensor* mask_list[4]; + for (int m = 0; m < num_mask_tokens; ++m) { + // Extract token m: [D, 1, 1] + auto* tok = ggml_view_3d(ctx, mask_tokens_out, D, 1, 1, + mask_tokens_out->nb[1], mask_tokens_out->nb[2], + m * mask_tokens_out->nb[1]); + tok = ggml_cont(ctx, tok); // [D, 1, 1] - int64_t numel = 1; - std::string line; - std::getline(shape_f, line); - size_t pos = 0; - while (pos < line.size()) { - size_t end = line.find(',', pos); - if (end == std::string::npos) { - end = line.size(); - } - if (end > pos) { - numel *= std::stoll(line.substr(pos, end - pos)); - } - pos = end + 1; - } + // MLP: 3 layers, 256→256→256→32, ReLU on first two + auto* hyper = sam3_mlp_forward(ctx, tok, + dec.hyper_w[m], dec.hyper_b[m], 3); + // hyper: [32, 1, 1] - if (expected_numel >= 0 && numel != expected_numel) { - fprintf(stderr, "%s: %s expected %d elements, got %lld\n", - __func__, path.c_str(), expected_numel, (long long)numel); - return false; + // Dot product: hyper^T @ up_flat → [1, 288*288, 1] + // Use mul_mat: up_flat^T [288*288, 32] @ hyper [32, 1] → [288*288, 1, 1] + auto* mask = ggml_mul_mat(ctx, up_flat, hyper); // [288*288, 1, 1] + mask_list[m] = mask; } - std::ifstream data_f(path + ".bin", std::ios::binary); - if (!data_f) { - fprintf(stderr, "%s: missing %s.bin\n", __func__, path.c_str()); - return false; + // Stack masks: [288*288, 4, 1] + auto* masks = mask_list[0]; + for (int m = 1; m < num_mask_tokens; ++m) { + masks = ggml_concat(ctx, masks, mask_list[m], 1); } + ggml_set_name(masks, "sam_dec_masks"); - data.resize((size_t)numel); - data_f.read(reinterpret_cast(data.data()), numel * sizeof(float)); - return data_f.good() || data_f.eof(); -} + // ── IoU prediction ─────────────────────────────────────────────────── + // iou_token_out: [D, 1, 1] + auto* iou_pred = sam3_mlp_forward(ctx, iou_token_out, + dec.iou_head_w, dec.iou_head_b, 3, + /*sigmoid_output=*/true); + // iou_pred: [4, 1, 1] → reshape to [4, 1] + iou_pred = ggml_reshape_2d(ctx, iou_pred, num_mask_tokens, 1); + ggml_set_name(iou_pred, "sam_dec_iou"); -static std::vector sam3_reorder_nchw_to_ggml_dwh(const std::vector& src, - int channels, - int height, - int width) { - std::vector dst((size_t)channels * height * width); - for (int c = 0; c < channels; ++c) { - for (int y = 0; y < height; ++y) { - for (int x = 0; x < width; ++x) { - const size_t src_idx = (size_t)c * height * width + (size_t)y * width + x; - const size_t dst_idx = (size_t)c + (size_t)x * channels + (size_t)y * channels * width; - dst[dst_idx] = src[src_idx]; - } - } + // ── Object score ───────────────────────────────────────────────────── + struct ggml_tensor* obj_score; + if (has_obj_score) { + // obj_in: [D, 1, 1] → MLP → [1, 1, 1] + obj_score = sam3_mlp_forward(ctx, obj_in, + dec.obj_head_w, dec.obj_head_b, 3); + obj_score = ggml_reshape_2d(ctx, obj_score, 1, 1); + } else { + // No obj_score prediction — return raw logit 10.0 (sigmoid ≈ 1.0, object always present) + obj_score = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, 1); + ggml_set_name(obj_score, "sam_dec_obj_score"); + ggml_set_input(obj_score); + // Mark that callers must set this to 10.0f before compute } - return dst; -} -static bool sam3_load_kv_text_file(const std::string& path, - std::map& kv) { - std::ifstream f(path); - if (!f) { - fprintf(stderr, "%s: missing %s\n", __func__, path.c_str()); - return false; - } + sam3_dec_result res; + res.masks = masks; + res.iou_pred = iou_pred; + res.obj_score = obj_score; + res.sam_token = sam_token; + res.mask_tokens = mask_tokens_out; + return res; +} - std::string line; - while (std::getline(f, line)) { - if (line.empty()) { - continue; - } - const size_t eq = line.find('='); - if (eq == std::string::npos) { - continue; - } - kv[line.substr(0, eq)] = line.substr(eq + 1); - } +/***************************************************************************** +** Image segmentation — PVS (visual-prompted) (Phase 6, Step 6.3) +*****************************************************************************/ - return true; -} +sam3_result sam3_segment_pvs(sam3_state& state, + const sam3_model& model, + const sam3_pvs_params& params) { +#if SAM3_LOG_LEVEL >= 1 + auto t_start = std::chrono::high_resolution_clock::now(); +#endif + const auto& hp = model.hparams; + const int D = hp.sam_embed_dim; // 256 + const int H = sam3_eff_feat_size(state, hp); + const int num_mask_tokens = hp.sam_n_multimask + 1; // 4 + const int eff_img_size = sam3_eff_img_size(state, hp); + sam3_result result; -static int sam3_meta_get_int(const std::map& kv, - const std::string& key, - int default_value = 0) { - auto it = kv.find(key); - if (it == kv.end()) { - return default_value; + // ── Validate ───────────────────────────────────────────────────────── + if (!state.neck_trk[0]) { + fprintf(stderr, "%s: image not encoded — call sam3_encode_image first\n", __func__); + return result; } - return std::stoi(it->second); -} - -bool sam3_test_dump_text_encoder(const sam3_model& model, - const std::vector& token_ids, - const std::string& output_dir, - int n_threads) { - const int L = model.hparams.text_ctx_len; - if ((int)token_ids.size() != L) { - fprintf(stderr, "%s: expected %d token IDs, got %zu\n", - __func__, L, token_ids.size()); - return false; + if (params.pos_points.empty() && !params.use_box) { + fprintf(stderr, "%s: no prompts provided (need at least one point or box)\n", __func__); + return result; } + SAM3_LOG(2, "%s: %zu pos points, %zu neg points, box=%s, multimask=%s\n", + __func__, params.pos_points.size(), params.neg_points.size(), + params.use_box ? "yes" : "no", params.multimask ? "yes" : "no"); + + // ── Build computation graph ────────────────────────────────────────── const size_t buf_size = ggml_tensor_overhead() * 8192 + ggml_graph_overhead() * 2; struct ggml_init_params gparams = { /*.mem_size =*/buf_size, @@ -12667,1773 +9523,1556 @@ bool sam3_test_dump_text_encoder(const sam3_model& model, struct ggml_context* ctx0 = ggml_init(gparams); if (!ctx0) { fprintf(stderr, "%s: failed to init compute context\n", __func__); - return false; + return result; } - auto* inp_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, L); - ggml_set_name(inp_tokens, "text_token_ids"); - ggml_set_input(inp_tokens); - - auto* text_features = sam3_build_text_encoder_graph(ctx0, inp_tokens, model); - std::vector tensor_names = { - "causal_mask", - "text_token_embed", - "text_after_pos_embed", - "text_final_ln", - "text_features_2d", - }; - for (int i = 0; i < model.hparams.text_layers; ++i) { - char name[64]; - - snprintf(name, sizeof(name), "text_block_%02d_after_ln1", i); - tensor_names.emplace_back(name); - - snprintf(name, sizeof(name), "text_block_%02d_qkv", i); - tensor_names.emplace_back(name); - - snprintf(name, sizeof(name), "text_block_%02d_attn_out", i); - tensor_names.emplace_back(name); - - snprintf(name, sizeof(name), "text_block_%02d_after_attn_residual", i); - tensor_names.emplace_back(name); + // ── SAM prompt encoder (CPU pre-compute + input tensors) ───────────── + auto pe_out = sam3_build_sam_pe(ctx0, params, D, H); - snprintf(name, sizeof(name), "text_block_%02d_after_ln2", i); - tensor_names.emplace_back(name); + // ── Create fresh input tensors for tracker features ────────────────── + // CRITICAL: Do NOT use state.neck_trk[*] directly in graph operations + // (like ggml_add). They are tensors from a previous graph, and + // ggml_build_forward_expand traces all ancestors — which would pull in + // the ENTIRE ViT + neck recomputation (2500+ nodes, ~37 seconds). + // Instead: create fresh input tensors and copy data from state on CPU. + const int H0 = H * 4; // 288 + const int H1 = H * 2; // 144 - snprintf(name, sizeof(name), "text_block_%02d_mlp_fc1", i); - tensor_names.emplace_back(name); + auto* image_feats = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, D, H, H, 1); + ggml_set_name(image_feats, "sam_dec_image_feats"); + ggml_set_input(image_feats); - snprintf(name, sizeof(name), "text_block_%02d_mlp_gelu", i); - tensor_names.emplace_back(name); + auto* feat_s0 = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, D, H0, H0, 1); + ggml_set_name(feat_s0, "pvs_feat_s0"); + ggml_set_input(feat_s0); - snprintf(name, sizeof(name), "text_block_%02d_mlp_out", i); - tensor_names.emplace_back(name); + auto* feat_s1 = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, D, H1, H1, 1); + ggml_set_name(feat_s1, "pvs_feat_s1"); + ggml_set_input(feat_s1); - snprintf(name, sizeof(name), "text_block_%02d_out", i); - tensor_names.emplace_back(name); - } + // ── SAM mask decoder graph ─────────────────────────────────────────── + auto dec_out = sam3_build_sam_dec_graph(ctx0, model, + image_feats, + pe_out.image_pe, + pe_out.sparse, + pe_out.dense, + feat_s0, + feat_s1, H); - ggml_set_output(text_features); - for (const auto& name : tensor_names) { - auto* t = ggml_get_tensor(ctx0, name.c_str()); - if (t && t != inp_tokens) { - ggml_set_output(t); - } - } + // Mark outputs + ggml_set_output(dec_out.masks); + ggml_set_output(dec_out.iou_pred); + ggml_set_output(dec_out.obj_score); + ggml_set_output(dec_out.sam_token); - struct ggml_cgraph* graph = ggml_new_graph_custom(ctx0, 8192, false); - ggml_build_forward_expand(graph, text_features); + // ── Build and allocate graph ───────────────────────────────────────── + struct ggml_cgraph* graph = ggml_new_graph_custom(ctx0, 32768, false); + ggml_build_forward_expand(graph, dec_out.masks); + ggml_build_forward_expand(graph, dec_out.iou_pred); + ggml_build_forward_expand(graph, dec_out.obj_score); + ggml_build_forward_expand(graph, dec_out.sam_token); auto* galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); - if (!ggml_gallocr_reserve(galloc, graph) || !ggml_gallocr_alloc_graph(galloc, graph)) { - fprintf(stderr, "%s: failed to allocate text encoder graph\n", __func__); + if (!ggml_gallocr_reserve(galloc, graph)) { + fprintf(stderr, "%s: failed to reserve graph memory\n", __func__); ggml_gallocr_free(galloc); ggml_free(ctx0); - return false; + return result; } - - ggml_backend_tensor_set(inp_tokens, token_ids.data(), 0, L * sizeof(int32_t)); - - auto* causal_mask = ggml_get_tensor(ctx0, "causal_mask"); - if (!causal_mask) { - fprintf(stderr, "%s: causal_mask tensor not found\n", __func__); + if (!ggml_gallocr_alloc_graph(galloc, graph)) { + fprintf(stderr, "%s: failed to allocate graph\n", __func__); ggml_gallocr_free(galloc); ggml_free(ctx0); - return false; + return result; } - std::vector mask_data(L * L); - sam3_fill_causal_mask(mask_data.data(), L); - ggml_backend_tensor_set(causal_mask, mask_data.data(), 0, L * L * sizeof(ggml_fp16_t)); + SAM3_LOG(2, "%s: graph allocated, %d nodes\n", __func__, ggml_graph_n_nodes(graph)); - sam3_graph_compute(model.backend, graph, n_threads); + // Set default obj_score when pred_obj_scores=False (older SAM2 models) + if (!model.sam_dec.obj_score_token) { + auto* t = ggml_graph_get_tensor(graph, "sam_dec_obj_score"); + if (t) { float v = 10.0f; ggml_backend_tensor_set(t, &v, 0, sizeof(float)); } + } - bool ok = true; - for (const auto& name : tensor_names) { - auto* t = ggml_get_tensor(ctx0, name.c_str()); - if (!t) { - fprintf(stderr, "%s: tensor '%s' not found in graph\n", __func__, name.c_str()); - ok = false; - continue; + // ── Upload input data (using cached embeddings) ──────────────────── + // Populate PE cache on first call (reads model weights from GPU once) + sam3_populate_pe_cache(state, model); + + { + const int N_pts = pe_out.n_tokens; + const int num_pos_feats = D / 2; + + std::vector all_coords; + std::vector all_labels; + sam3_collect_pvs_prompt_tokens(params, all_coords, all_labels); + + // Sparse embeddings — only this changes per call + std::vector sparse_data(N_pts * D, 0.0f); + for (int p = 0; p < N_pts; ++p) { + // Scale from original image space to model input space, then shift to pixel center + float px = all_coords[p * 2 + 0] / (float)state.orig_width * (float)eff_img_size + 0.5f; + float py = all_coords[p * 2 + 1] / (float)state.orig_height * (float)eff_img_size + 0.5f; + float x_norm = px / (float)eff_img_size; + float y_norm = py / (float)eff_img_size; + float pe_vec[256]; + sam3_pe_encode_coord(pe_vec, x_norm, y_norm, + state.pe_gauss_cache.data(), num_pos_feats); + int label = all_labels[p]; + if (label == -1) { + for (int d = 0; d < D; ++d) + sparse_data[p * D + d] = state.not_a_point_cache[d]; + } else { + for (int d = 0; d < D; ++d) + sparse_data[p * D + d] = pe_vec[d] + state.point_emb_cache[label][d]; + } } - if (!sam3_dump_tensor_to_path(t, name, output_dir + "/" + name)) { - ok = false; + ggml_backend_tensor_set(pe_out.sparse, sparse_data.data(), 0, N_pts * D * sizeof(float)); + + // Dump sparse embeddings if requested + { + const char* dd = getenv("SAM2_DUMP_DIR"); + if (dd) { + char p[512]; snprintf(p, sizeof(p), "%s/cpp_sparse_emb.bin", dd); + FILE* f = fopen(p, "wb"); + if (f) { fwrite(sparse_data.data(), sizeof(float), N_pts * D, f); fclose(f); } + fprintf(stderr, " [DUMP] cpp_sparse_emb: %d tokens x %d dims\n", N_pts, D); + } } + + // Dense PE grid and no-mask embedding — use pre-computed caches + ggml_backend_tensor_set(pe_out.image_pe, state.dense_pe_cache.data(), + 0, D * H * H * sizeof(float)); + ggml_backend_tensor_set(pe_out.dense, state.dense_nomask_cache.data(), + 0, D * H * H * sizeof(float)); } - ggml_gallocr_free(galloc); - ggml_free(ctx0); - return ok; -} + // ── Copy tracker features from state to fresh input tensors ───────── + // image_feats = neck_trk[2] + no_mem_embed (computed on CPU) + { + const int n2 = D * H * H; + std::vector trk2(n2), no_mem_data(D); + ggml_backend_tensor_get(state.neck_trk[2], trk2.data(), 0, n2 * sizeof(float)); + ggml_backend_tensor_get(model.tensors.at("no_mem_embed"), no_mem_data.data(), 0, D * sizeof(float)); + // Add no_mem_embed (broadcast [D] to [D, H, H]) + for (int s = 0; s < H * H; ++s) + for (int d = 0; d < D; ++d) + trk2[d + s * D] += no_mem_data[d]; + ggml_backend_tensor_set(image_feats, trk2.data(), 0, n2 * sizeof(float)); -bool sam3_test_dump_phase5(const sam3_model& model, - const sam3_state& state, - const std::vector& token_ids, - const std::string& output_dir, - int n_threads) { - const auto& hp = model.hparams; - const int D = hp.neck_dim; - const int H = hp.n_img_embd(); - const int L = hp.text_ctx_len; - const int NQ = hp.ddec_num_queries; + // Dump image_feats (with no_mem_embed) if requested + { + const char* dd = getenv("SAM2_DUMP_DIR"); + if (dd) { + char p[512]; snprintf(p, sizeof(p), "%s/cpp_image_feats.bin", dd); + FILE* f = fopen(p, "wb"); + if (f) { fwrite(trk2.data(), sizeof(float), n2, f); fclose(f); } + fprintf(stderr, " [DUMP] cpp_image_feats: [%d, %d, %d]\n", D, H, H); + } + } - if ((int)token_ids.size() != L) { - fprintf(stderr, "%s: expected %d token IDs, got %zu\n", - __func__, L, token_ids.size()); - return false; + // feat_s0 = neck_trk[0], feat_s1 = neck_trk[1] + const int n0 = D * H0 * H0; + const int n1 = D * H1 * H1; + std::vector trk0(n0), trk1(n1); + ggml_backend_tensor_get(state.neck_trk[0], trk0.data(), 0, n0 * sizeof(float)); + ggml_backend_tensor_set(feat_s0, trk0.data(), 0, n0 * sizeof(float)); + ggml_backend_tensor_get(state.neck_trk[1], trk1.data(), 0, n1 * sizeof(float)); + ggml_backend_tensor_set(feat_s1, trk1.data(), 0, n1 * sizeof(float)); } - if (!state.neck_det[0] || !state.neck_det_pe[2]) { - fprintf(stderr, "%s: encoded detector features are missing\n", __func__); - return false; + + // ── Compute ────────────────────────────────────────────────────────── + { +#if SAM3_LOG_LEVEL >= 1 + auto t0 = std::chrono::high_resolution_clock::now(); +#endif + if (!sam3_graph_compute(model.backend, graph, state.n_threads)) { + ggml_gallocr_free(galloc); + ggml_free(ctx0); + return result; + } +#if SAM3_LOG_LEVEL >= 1 + auto t1 = std::chrono::high_resolution_clock::now(); + double ms = std::chrono::duration(t1 - t0).count(); + SAM3_LOG(1, "%s: graph computed in %.1f ms (%d threads)\n", + __func__, ms, state.n_threads); +#endif + } + + // ── Dump decoder outputs if SAM2_DUMP_DIR set ────────────────────── + { + const char* dump_dir = getenv("SAM2_DUMP_DIR"); + if (dump_dir) { + auto dump_t = [&](const char* name, struct ggml_tensor* t) { + if (!t) return; + int64_t nb = ggml_nbytes(t); + std::vector buf(nb); + ggml_backend_tensor_get(t, buf.data(), 0, nb); + char path[512]; + snprintf(path, sizeof(path), "%s/%s.bin", dump_dir, name); + FILE* f = fopen(path, "wb"); + if (f) { fwrite(buf.data(), 1, nb, f); fclose(f); } + snprintf(path, sizeof(path), "%s/%s.shape", dump_dir, name); + f = fopen(path, "w"); + if (f) { + fprintf(f, "%lld,%lld,%lld,%lld", + (long long)t->ne[0], (long long)t->ne[1], + (long long)t->ne[2], (long long)t->ne[3]); + fclose(f); + } + fprintf(stderr, " [DUMP] %s: [%lld,%lld,%lld,%lld]\n", name, + (long long)t->ne[0], (long long)t->ne[1], + (long long)t->ne[2], (long long)t->ne[3]); + }; + dump_t("cpp_pvs_masks", dec_out.masks); + dump_t("cpp_pvs_iou", dec_out.iou_pred); + dump_t("cpp_pvs_obj_score", dec_out.obj_score); + // Decoder transformer intermediates + dump_t("cpp_dec_queries", ggml_graph_get_tensor(graph, "dbg_dec_queries_out")); + dump_t("cpp_dec_keys", ggml_graph_get_tensor(graph, "dbg_dec_keys_out")); + // Block 0 internals + const char* b0_names[] = {"sam_dec_tokens_initial", + "dbg_twoway_skip_sa_norm", "dbg_twoway_skip_ca_tok2img", + "dbg_twoway_skip_mlp", "dbg_twoway_skip_img2tok", + "dbg_sa0_Q_proj", "dbg_sa0_merged"}; + for (auto* bn : b0_names) dump_t(bn, ggml_graph_get_tensor(graph, bn)); + // Per-block outputs + for (int bi = 0; bi < 2; bi++) { + char bn[64]; + snprintf(bn, sizeof(bn), "sam_dec_block%d_queries", bi); + dump_t(bn, ggml_graph_get_tensor(graph, bn)); + snprintf(bn, sizeof(bn), "sam_dec_block%d_keys", bi); + dump_t(bn, ggml_graph_get_tensor(graph, bn)); + } + } } - const size_t buf_size = ggml_tensor_overhead() * 65536 + ggml_graph_overhead() * 2; - struct ggml_init_params gparams = { - /*.mem_size =*/buf_size, - /*.mem_buffer =*/nullptr, - /*.no_alloc =*/true, - }; - struct ggml_context* ctx0 = ggml_init(gparams); - if (!ctx0) { - fprintf(stderr, "%s: failed to init compute context\n", __func__); - return false; - } + // ── Read outputs ───────────────────────────────────────────────────── + // masks: [H*4×H*4, 4, 1] + const int mask_hw = H * 4; + std::vector masks_data(mask_hw * mask_hw * num_mask_tokens); + ggml_backend_tensor_get(dec_out.masks, masks_data.data(), 0, masks_data.size() * sizeof(float)); + + // IoU predictions: [4, 1] + std::vector iou_data(num_mask_tokens); + ggml_backend_tensor_get(dec_out.iou_pred, iou_data.data(), 0, num_mask_tokens * sizeof(float)); + + // Object score: [1, 1] + float obj_logit = 0.0f; + ggml_backend_tensor_get(dec_out.obj_score, &obj_logit, 0, sizeof(float)); + float obj_score = 1.0f / (1.0f + expf(-obj_logit)); + + // SAM output token: [D, 1] — needed for object pointer extraction in tracking + std::vector sam_token_data(D); + ggml_backend_tensor_get(dec_out.sam_token, sam_token_data.data(), 0, D * sizeof(float)); + + SAM3_LOG(2, "%s: obj_score=%.4f (logit=%.4f), iou=[%.3f, %.3f, %.3f, %.3f]\n", + __func__, obj_score, obj_logit, + iou_data[0], iou_data[1], iou_data[2], iou_data[3]); - auto* inp_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, L); - ggml_set_name(inp_tokens, "text_token_ids"); - ggml_set_input(inp_tokens); + // ── Select masks based on multimask mode ───────────────────────────── + // Python: if multimask_output → masks[:, 1:, :, :], iou_pred[:, 1:] + // else → masks[:, 0:1, :, :], iou_pred[:, 0:1] + int start_idx, end_idx; + if (params.multimask) { + start_idx = 1; + end_idx = num_mask_tokens; + } else { + start_idx = 0; + end_idx = 1; + } - auto* text_features_2d = sam3_build_text_encoder_graph(ctx0, inp_tokens, model); - auto* text_features = ggml_reshape_3d(ctx0, text_features_2d, D, L, 1); - ggml_set_name(text_features, "text_features"); + for (int m = start_idx; m < end_idx; ++m) { + sam3_detection det; + det.sam_token = sam_token_data; - // Make a snapshot copy of text_features for dumping — the graph allocator - // may reuse the view's underlying buffer for later ops. - auto* text_features_snap = ggml_cont(ctx0, ggml_reshape_2d(ctx0, text_features_2d, D, L)); - ggml_set_name(text_features_snap, "text_features_snap"); + // Resize mask from 288×288 to original image size + const float* mask_ptr = masks_data.data() + m * mask_hw * mask_hw; + auto mask_resized = sam3_bilinear_interpolate(mask_ptr, mask_hw, mask_hw, + state.orig_width, state.orig_height); - auto* img_feats = ggml_reshape_3d(ctx0, state.neck_det[2], D, H * H, 1); - auto* img_pe = ggml_reshape_3d(ctx0, state.neck_det_pe[2], D, H * H, 1); + // Binarize at threshold 0.0 (sigmoid(logit) > 0.5 ↔ logit > 0.0) + det.mask.width = state.orig_width; + det.mask.height = state.orig_height; + det.mask.data.resize(state.orig_width * state.orig_height); + for (int i = 0; i < (int)mask_resized.size(); ++i) { + det.mask.data[i] = (mask_resized[i] > 0.0f) ? 255 : 0; + } - auto* sine_dim_t = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 1, 64); - ggml_set_name(sine_dim_t, "sine_dim_t"); - ggml_set_input(sine_dim_t); + det.mask.iou_score = iou_data[m]; + det.mask.obj_score = obj_score; + det.mask.instance_id = m; + det.score = iou_data[m]; + det.iou_score = iou_data[m]; + det.instance_id = m; - auto* rpb_coords = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, H); - ggml_set_name(rpb_coords, "rpb_coords"); - ggml_set_input(rpb_coords); + // Compute bounding box from mask + int min_x = state.orig_width, min_y = state.orig_height; + int max_x = 0, max_y = 0; + for (int y = 0; y < state.orig_height; ++y) { + for (int x = 0; x < state.orig_width; ++x) { + if (det.mask.data[y * state.orig_width + x] > 0) { + min_x = std::min(min_x, x); + min_y = std::min(min_y, y); + max_x = std::max(max_x, x); + max_y = std::max(max_y, y); + } + } + } + det.box = {(float)min_x, (float)min_y, (float)max_x, (float)max_y}; - auto* text_valid_mask = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, L, 1, 1); - ggml_set_name(text_valid_mask, "text_valid_mask"); - ggml_set_input(text_valid_mask); + result.detections.push_back(std::move(det)); + } - auto* text_attn_bias = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, L, 1, 1); - ggml_set_name(text_attn_bias, "text_attn_bias"); - ggml_set_input(text_attn_bias); + SAM3_LOG(2, "%s: %zu masks returned\n", __func__, result.detections.size()); - auto* conditioned = sam3_build_fenc_graph(ctx0, model, img_feats, text_features, - img_pe, text_attn_bias); - auto* fenc_layer5_out = ggml_cont(ctx0, conditioned); - ggml_set_name(fenc_layer5_out, "fenc_layer5_out"); - ggml_set_name(conditioned, "fenc_output"); + // ── Cleanup ────────────────────────────────────────────────────────── + ggml_gallocr_free(galloc); + ggml_free(ctx0); - auto ddec_out = sam3_build_ddec_graph(ctx0, model, conditioned, img_pe, text_features, - sine_dim_t, rpb_coords, text_attn_bias, text_valid_mask); +#if SAM3_LOG_LEVEL >= 1 + auto t_end = std::chrono::high_resolution_clock::now(); + double total_ms = std::chrono::duration(t_end - t_start).count(); + SAM3_LOG(1, "%s: completed in %.1f ms\n", __func__, total_ms); +#endif - struct ggml_tensor* fpn_feats[3] = { - state.neck_det[0], - state.neck_det[1], - state.neck_det[2], - }; + return result; +} - auto* obj_queries = ggml_view_3d(ctx0, ddec_out.queries, D, NQ, 1, - ddec_out.queries->nb[1], ddec_out.queries->nb[2], - 1 * ddec_out.queries->nb[1]); - obj_queries = ggml_cont(ctx0, obj_queries); +/***************************************************************************** +** Video tracking (Phase 7) +*****************************************************************************/ - auto* mask_logits = sam3_build_seg_head_graph(ctx0, model, conditioned, fpn_feats, - obj_queries, text_features, text_attn_bias); - ggml_set_name(mask_logits, "seg_mask_logits"); +struct sam3_prop_output { + std::vector mask_logits; + std::vector iou_scores; + float obj_score; + std::vector sam_token; + int n_masks, mask_h, mask_w; +}; - struct named_tensor { - const char* name; - struct ggml_tensor* tensor; - }; +// Lazily compute and cache PE/RoPE data that is identical across all propagation calls. +static void sam3_ensure_tracker_pe_caches(sam3_tracker& tracker, const sam3_hparams& hp, + int eff_feat_size = 0) { + const int H = (eff_feat_size > 0) ? eff_feat_size : hp.feat_size(); + if (tracker.pe_caches_valid) return; - const std::vector direct_tensors = { - {"text_features", text_features_snap}, - {"text_valid_mask", text_valid_mask}, - {"fenc_img_input", img_feats}, - {"fenc_pos_embed", img_pe}, - {"fenc_prompt", text_features_snap}, - {"img_pe_72", state.neck_det_pe[2]}, - {"ddec_query_embed", model.ddec.query_embed}, - {"ddec_ref_pts_raw", model.tensors.at("ddec.reference_points.weight")}, - {"ddec_presence_token", model.ddec.presence_token}, - {"ddec_pred_boxes", ddec_out.pred_boxes}, - {"ddec_presence_logit", ddec_out.presence_score}, - }; + const int D = hp.neck_dim; // 256 + const int MD = hp.mem_out_dim; // 64 + const int N = H * H; + const int half_d = D / 2; // 128 - std::vector named_outputs = { - "fenc_output", - "ddec_ref_boxes_init", - "ddec_query_sine_0", - "ddec_query_pos_0", - "ddec_rpb_mask_0", - "ddec_layer0_after_sa", - "ddec_layer0_after_text_ca", - "ddec_layer0_after_img_ca", - "ddec_layer0_full_out", - "ddec_layer0_presence", - "ddec_normed_output", - "scoring_prompt_mlp_out", - "scoring_pooled", - "scoring_proj_pooled", - "scoring_proj_hs", - "scoring_class_scores", - "seg_enc_after_ca", - "seg_enc_visual", - "seg_pixel_dec_stage0", - "seg_pixel_dec_stage1", - "seg_pixel_decoder_out", - "seg_instance_embed", - "seg_mask_embed", - "seg_mask_logits", - }; + tracker.cached_sinpe_256 = sam3_sinusoidal_pe_2d(H, H, D); + tracker.cached_sinpe_64 = sam3_sinusoidal_pe_2d(H, H, MD); - for (int i = 0; i < hp.fenc_layers; ++i) { - char name[64]; - snprintf(name, sizeof(name), "fenc_layer%d_out", i); - named_outputs.emplace_back(name); - } - for (int i = 0; i < hp.ddec_layers; ++i) { - char out_name[64]; - char box_name[64]; - snprintf(out_name, sizeof(out_name), "ddec_layer%d_out", i); - snprintf(box_name, sizeof(box_name), "ddec_layer%d_refboxes", i); - named_outputs.emplace_back(out_name); - named_outputs.emplace_back(box_name); - } - - ggml_set_output(text_features_snap); - ggml_set_output(text_valid_mask); - ggml_set_output(img_feats); - ggml_set_output(img_pe); - for (const auto& name : named_outputs) { - auto* t = ggml_get_tensor(ctx0, name.c_str()); - if (t) { - ggml_set_output(t); - } - } - ggml_set_output(mask_logits); - ggml_set_output(ddec_out.class_scores); - ggml_set_output(ddec_out.pred_boxes); - ggml_set_output(ddec_out.presence_score); - - struct ggml_cgraph* graph = ggml_new_graph_custom(ctx0, 65536, false); - ggml_build_forward_expand(graph, ddec_out.class_scores); - ggml_build_forward_expand(graph, ddec_out.pred_boxes); - ggml_build_forward_expand(graph, ddec_out.presence_score); - ggml_build_forward_expand(graph, mask_logits); - ggml_build_forward_expand(graph, text_features_snap); - ggml_build_forward_expand(graph, text_valid_mask); - ggml_build_forward_expand(graph, img_feats); - ggml_build_forward_expand(graph, img_pe); - for (const auto& name : named_outputs) { - auto* t = ggml_get_tensor(ctx0, name.c_str()); - if (t) { - ggml_build_forward_expand(graph, t); + // Compute axial CIS and reorder to [2, half_d, N] layout + std::vector rope_raw(N * D); + sam3_compute_axial_cis(rope_raw.data(), D, H, H, 10000.0f, 1.0f); + tracker.cached_axial_cis_reord.resize(2 * half_d * N); + for (int n = 0; n < N; ++n) + for (int i = 0; i < half_d; ++i) { + tracker.cached_axial_cis_reord[0 + i * 2 + n * D] = rope_raw[n * D + i * 2 + 0]; + tracker.cached_axial_cis_reord[1 + i * 2 + n * D] = rope_raw[n * D + i * 2 + 1]; } - } - auto* galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); - if (!ggml_gallocr_reserve(galloc, graph) || !ggml_gallocr_alloc_graph(galloc, graph)) { - fprintf(stderr, "%s: failed to allocate phase 5 graph\n", __func__); - ggml_gallocr_free(galloc); - ggml_free(ctx0); - return false; + // EdgeTAM: compute 16x16 RoPE for cross-attention K (perceiver 2D latents) + if (hp.has_perceiver) { + const int K16 = 16; // perceiver 2D grid size + const int NK = K16 * K16; // 256 tokens + std::vector rope_k16_raw(NK * D); + sam3_compute_axial_cis(rope_k16_raw.data(), D, K16, K16, 10000.0f, 1.0f); + tracker.cached_axial_cis_k16_reord.resize(2 * half_d * NK); + for (int n = 0; n < NK; ++n) + for (int i = 0; i < half_d; ++i) { + tracker.cached_axial_cis_k16_reord[0 + i * 2 + n * D] = rope_k16_raw[n * D + i * 2 + 0]; + tracker.cached_axial_cis_k16_reord[1 + i * 2 + n * D] = rope_k16_raw[n * D + i * 2 + 1]; + } + SAM3_LOG(2, "%s: EdgeTAM K16 RoPE cache: %zu floats\n", __func__, + tracker.cached_axial_cis_k16_reord.size()); } - ggml_backend_tensor_set(inp_tokens, token_ids.data(), 0, L * sizeof(int32_t)); + tracker.pe_caches_valid = true; + SAM3_LOG(2, "%s: tracker PE caches populated (%.1f KB)\n", __func__, + (tracker.cached_sinpe_256.size() + tracker.cached_sinpe_64.size() + + tracker.cached_axial_cis_reord.size()) * sizeof(float) / 1024.0f); +} - auto* causal_mask = ggml_get_tensor(ctx0, "causal_mask"); - if (causal_mask) { - std::vector mask_data(L * L); - sam3_fill_causal_mask(mask_data.data(), L); - ggml_backend_tensor_set(causal_mask, mask_data.data(), 0, L * L * sizeof(ggml_fp16_t)); - } +static sam3_prop_output sam3_propagate_single( + sam3_tracker& tracker, sam3_state& state, const sam3_model& model, + const sam3_masklet& masklet, + const std::vector& mem_bank, + const std::vector>& ptr_bank) { + sam3_prop_output output = {}; + const auto& hp = model.hparams; + const int D = hp.neck_dim, MD = hp.mem_out_dim; + const int H = sam3_eff_feat_size(state, hp); + const int N = H * H; - { - std::vector dim_t_data(64); - for (int i = 0; i < 64; ++i) { - dim_t_data[i] = 2.0f * (float)M_PI / powf(10000.0f, 2.0f * (float)i / 128.0f); - } - ggml_backend_tensor_set(sine_dim_t, dim_t_data.data(), 0, 64 * sizeof(float)); - } + auto sel = sam3_select_memory_frames(mem_bank, hp.num_maskmem); + if (sel.empty()) return output; - { - std::vector coords(H); - for (int i = 0; i < H; ++i) { - coords[i] = (float)i / (float)H; + // ── Build prompt and prompt_pos via sam3_build_prompt_and_pos ───────── + // EdgeTAM with perceiver: each slot has 512 tokens (not H*H). + // Standard SAM2/SAM3: each slot has H*H tokens. + const bool use_perceiver = hp.has_perceiver != 0; + const int N_per_slot = use_perceiver ? (hp.perceiver_n_latents_1d + hp.perceiver_n_latents_2d) : N; + + int n_sel = (int)sel.size(); + std::vector> slot_feats(n_sel), slot_pes(n_sel); + std::vector spatial_tpos(n_sel, 1); // default t_pos=1 for non-cond + for (int s = 0; s < n_sel; ++s) { + slot_feats[s].resize(MD * N_per_slot); + ggml_backend_tensor_get(mem_bank[sel[s]].spatial_feats, + slot_feats[s].data(), 0, MD * N_per_slot * sizeof(float)); + slot_pes[s].resize(MD * N_per_slot); + if (mem_bank[sel[s]].spatial_pe) { + ggml_backend_tensor_get(mem_bank[sel[s]].spatial_pe, + slot_pes[s].data(), 0, MD * N_per_slot * sizeof(float)); + } else { + sam3_ensure_tracker_pe_caches(tracker, hp, H); + if (use_perceiver) { + // For perceiver: zeros for 1D tokens, sinusoidal for 2D tokens + const int N_1d = hp.perceiver_n_latents_1d; + const int N_2d = hp.perceiver_n_latents_2d; + memset(slot_pes[s].data(), 0, MD * N_1d * sizeof(float)); + auto pe_2d = sam3_sinusoidal_pe_2d(16, 16, MD); + memcpy(slot_pes[s].data() + MD * N_1d, pe_2d.data(), MD * N_2d * sizeof(float)); + } else { + slot_pes[s] = tracker.cached_sinpe_64; + } } - ggml_backend_tensor_set(rpb_coords, coords.data(), 0, H * sizeof(float)); + spatial_tpos[s] = mem_bank[sel[s]].is_cond_frame ? 0 : (n_sel - s); } - auto* qpos_pres = ggml_get_tensor(ctx0, "ddec_query_pos_pres"); - if (qpos_pres) { - std::vector zeros(D, 0.0f); - ggml_backend_tensor_set(qpos_pres, zeros.data(), 0, D * sizeof(float)); + int P = std::min((int)ptr_bank.size(), hp.max_obj_ptrs); + std::vector> obj_ptrs(P); + std::vector ptr_tpos(P); + int cur_frame = tracker.frame_index; + for (int p = 0; p < P; ++p) { + obj_ptrs[p].resize(D); + ggml_backend_tensor_get(ptr_bank[p].second, obj_ptrs[p].data(), 0, D * sizeof(float)); + // Use actual frame distance (matches Python: abs(frame_idx - t)) + ptr_tpos[p] = std::abs(cur_frame - ptr_bank[p].first); + if (ptr_tpos[p] < 1) ptr_tpos[p] = 1; // minimum distance of 1 } - auto* rpb_pz = ggml_get_tensor(ctx0, "rpb_pres_zeros"); - if (rpb_pz) { - int n = (int)(rpb_pz->ne[0] * rpb_pz->ne[1] * rpb_pz->ne[2] * rpb_pz->ne[3]); - std::vector zeros(n, 0.0f); - ggml_backend_tensor_set(rpb_pz, zeros.data(), 0, n * sizeof(float)); - } + auto pd = sam3_build_prompt_and_pos(model, slot_feats, slot_pes, spatial_tpos, obj_ptrs, ptr_tpos, H); - { - int n_valid = 0; - for (int i = 0; i < L; ++i) { - if (token_ids[i] != 0) { - ++n_valid; + // ── RoPE frequencies (cached) ────────────────────────────────────── + sam3_ensure_tracker_pe_caches(tracker, hp, H); + const int half_d = D / 2; // 128 + const auto& rope_q_reord = tracker.cached_axial_cis_reord; + // For cross-attn K: build rope_k_data for all M_spatial tokens + std::vector rope_k_data; + if (pd.M_spatial > 0) { + rope_k_data.resize(2 * half_d * pd.M_spatial); + if (use_perceiver) { + // EdgeTAM perceiver: each frame has N_per_slot=512 tokens. + // First 256 (1D latents): identity RoPE (cos=1, sin=0). + // Last 256 (2D latents): 16x16 RoPE from cached_axial_cis_k16_reord. + const int N_1d = hp.perceiver_n_latents_1d; // 256 + const int N_2d = hp.perceiver_n_latents_2d; // 256 + const auto& rope_k16 = tracker.cached_axial_cis_k16_reord; // [2, 128, 256] + for (int s = 0; s < n_sel; ++s) { + float* dst = rope_k_data.data() + s * D * N_per_slot; + // 1D tokens: identity RoPE (cos=1, sin=0) in [2, half_d, N_1d] layout + // Layout: for token n, dim i: cos at [0 + i*2 + n*D], sin at [1 + i*2 + n*D] + for (int n = 0; n < N_1d; ++n) + for (int i = 0; i < half_d; ++i) { + dst[0 + i * 2 + n * D] = 1.0f; // cos = 1 + dst[1 + i * 2 + n * D] = 0.0f; // sin = 0 + } + // 2D tokens: copy 16x16 RoPE + float* dst_2d = dst + D * N_1d; + memcpy(dst_2d, rope_k16.data(), D * N_2d * sizeof(float)); } + } else { + // Standard: repeat HxH axial CIS for each memory frame + for (int s = 0; s < pd.M_spatial / N; ++s) + memcpy(rope_k_data.data() + s * D * N, rope_q_reord.data(), D * N * sizeof(float)); } - if (n_valid == 0) { - n_valid = 1; - } - - const float scale = (float)L / (float)n_valid; - std::vector valid_mask(L); - std::vector attn_bias(L); - for (int i = 0; i < L; ++i) { - const bool is_valid = token_ids[i] != 0; - valid_mask[i] = is_valid ? scale : 0.0f; - attn_bias[i] = is_valid ? 0.0f : -1.0e9f; - } - ggml_backend_tensor_set(text_valid_mask, valid_mask.data(), 0, L * sizeof(float)); - ggml_backend_tensor_set(text_attn_bias, attn_bias.data(), 0, L * sizeof(float)); } - sam3_graph_compute(model.backend, graph, n_threads); + // ── Build graph ───────────────────────────────────────────────────── + const size_t buf_size = ggml_tensor_overhead() * 32768 + ggml_graph_overhead() * 2; + struct ggml_init_params gparams = {buf_size, nullptr, true}; + auto* ctx0 = ggml_init(gparams); + if (!ctx0) return output; - bool ok = true; - for (const auto& item : direct_tensors) { - if (!sam3_dump_tensor_to_path(item.tensor, item.name, output_dir + "/" + item.name)) { - ok = false; - } - } - for (const auto& name : named_outputs) { - auto* t = ggml_get_tensor(ctx0, name.c_str()); - if (!t) { - fprintf(stderr, "%s: tensor '%s' not found in graph\n", __func__, name.c_str()); - ok = false; - continue; - } - if (!sam3_dump_tensor_to_path(t, name, output_dir + "/" + name)) { - ok = false; - } - } + // CRITICAL: create fresh input tensors for state features. + // Using state.neck_trk[*] directly as ggml_reshape operands pulls in + // the entire ViT+neck recomputation from the image encoder graph. + auto* curr = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, D, N, 1); + ggml_set_name(curr, "prop_curr"); + ggml_set_input(curr); - ggml_gallocr_free(galloc); - ggml_free(ctx0); - return ok; -} + // src_pos (sinusoidal PE 256-dim for 72×72) + auto* src_pos_t = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, D, N, 1); + ggml_set_name(src_pos_t, "src_pos"); + ggml_set_input(src_pos_t); -bool sam3_test_dump_phase5_from_ref_inputs(const sam3_model& model, - const std::vector& token_ids, - const std::string& prephase_ref_dir, - const std::string& phase5_ref_dir, - const std::string& output_dir, - int n_threads) { - const auto& hp = model.hparams; - const int D = hp.neck_dim; - const int H = hp.n_img_embd(); - const int L = hp.text_ctx_len; - const int NQ = hp.ddec_num_queries; - const int H1 = H * 2; - const int H0 = H * 4; - - if ((int)token_ids.size() != L) { - fprintf(stderr, "%s: expected %d token IDs, got %zu\n", - __func__, L, token_ids.size()); - return false; - } + // Prompt and prompt_pos + auto* prompt_t = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, MD, pd.M_total, 1); + ggml_set_name(prompt_t, "prompt"); + ggml_set_input(prompt_t); + auto* prompt_pos_t = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, MD, pd.M_total, 1); + ggml_set_name(prompt_pos_t, "prompt_pos"); + ggml_set_input(prompt_pos_t); - std::vector neck_det_0; - std::vector neck_det_1; - std::vector fenc_img_input_data; - std::vector fenc_pos_embed_data; - std::vector img_pe_72_data; - std::vector text_features_data; - if (!sam3_load_ref_f32_data(prephase_ref_dir + "/neck_det_0", neck_det_0, D * H0 * H0) || - !sam3_load_ref_f32_data(prephase_ref_dir + "/neck_det_1", neck_det_1, D * H1 * H1) || - !sam3_load_ref_f32_data(phase5_ref_dir + "/fenc_img_input", fenc_img_input_data, D * H * H) || - !sam3_load_ref_f32_data(phase5_ref_dir + "/fenc_pos_embed", fenc_pos_embed_data, D * H * H) || - !sam3_load_ref_f32_data(phase5_ref_dir + "/img_pe_72", img_pe_72_data, D * H * H) || - !sam3_load_ref_f32_data(phase5_ref_dir + "/text_features", text_features_data, D * L)) { - return false; + // RoPE frequencies + auto* rope_q_t = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 2, half_d, N); + ggml_set_name(rope_q_t, "rope_q"); + ggml_set_input(rope_q_t); + struct ggml_tensor* rope_k_t = nullptr; + if (pd.M_spatial > 0) { + rope_k_t = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 2, half_d, pd.M_spatial); + ggml_set_name(rope_k_t, "rope_k"); + ggml_set_input(rope_k_t); } - neck_det_0 = sam3_reorder_nchw_to_ggml_dwh(neck_det_0, D, H0, H0); - neck_det_1 = sam3_reorder_nchw_to_ggml_dwh(neck_det_1, D, H1, H1); + auto* conditioned = sam3_build_mem_attn_graph(ctx0, model, curr, src_pos_t, + prompt_t, prompt_pos_t, + rope_q_t, rope_k_t, + pd.num_obj_ptr_tokens); + auto* cond_spatial = ggml_reshape_4d(ctx0, conditioned, D, H, H, 1); - const size_t buf_size = ggml_tensor_overhead() * 65536 + ggml_graph_overhead() * 2; - struct ggml_init_params gparams = { - /*.mem_size =*/buf_size, - /*.mem_buffer =*/nullptr, - /*.no_alloc =*/true, - }; - struct ggml_context* ctx0 = ggml_init(gparams); - if (!ctx0) { - fprintf(stderr, "%s: failed to init compute context\n", __func__); - return false; - } + // Bug 3 fix: single not_a_point_embed token instead of empty sparse + auto* sparse_in = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, D, 1, 1); + ggml_set_name(sparse_in, "prop_sparse"); + ggml_set_input(sparse_in); - auto* text_features = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, D, L, 1); - ggml_set_name(text_features, "text_features"); - ggml_set_input(text_features); - - auto* neck0 = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, D, H0, H0, 1); - ggml_set_name(neck0, "ref_neck_det_0"); - ggml_set_input(neck0); - auto* neck1 = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, D, H1, H1, 1); - ggml_set_name(neck1, "ref_neck_det_1"); - ggml_set_input(neck1); - auto* img_feats = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, D, H * H, 1); - ggml_set_name(img_feats, "ref_fenc_img_input"); - ggml_set_input(img_feats); - auto* img_pe = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, D, H * H, 1); - ggml_set_name(img_pe, "ref_fenc_pos_embed"); - ggml_set_input(img_pe); - auto* img_pe_72_cmp = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, D * H * H); - ggml_set_name(img_pe_72_cmp, "ref_img_pe_72_flat"); - ggml_set_input(img_pe_72_cmp); - - auto* sine_dim_t = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 1, 64); - ggml_set_name(sine_dim_t, "sine_dim_t"); - ggml_set_input(sine_dim_t); - - auto* rpb_coords = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, H); - ggml_set_name(rpb_coords, "rpb_coords"); - ggml_set_input(rpb_coords); - - auto* text_valid_mask = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, L, 1, 1); - ggml_set_name(text_valid_mask, "text_valid_mask"); - ggml_set_input(text_valid_mask); - - auto* text_attn_bias = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, L, 1, 1); - ggml_set_name(text_attn_bias, "text_attn_bias"); - ggml_set_input(text_attn_bias); - - auto* conditioned = sam3_build_fenc_graph(ctx0, model, img_feats, text_features, - img_pe, text_attn_bias); - auto* fenc_layer5_out = ggml_cont(ctx0, conditioned); - ggml_set_name(fenc_layer5_out, "fenc_layer5_out"); - ggml_set_name(conditioned, "fenc_output"); - - auto ddec_out = sam3_build_ddec_graph(ctx0, model, conditioned, img_pe, text_features, - sine_dim_t, rpb_coords, text_attn_bias, text_valid_mask); - - struct ggml_tensor* fpn_feats[3] = {neck0, neck1, nullptr}; - - auto* obj_queries = ggml_view_3d(ctx0, ddec_out.queries, D, NQ, 1, - ddec_out.queries->nb[1], ddec_out.queries->nb[2], - 1 * ddec_out.queries->nb[1]); - obj_queries = ggml_cont(ctx0, obj_queries); - - auto* mask_logits = sam3_build_seg_head_graph(ctx0, model, conditioned, fpn_feats, - obj_queries, text_features, text_attn_bias); - ggml_set_name(mask_logits, "seg_mask_logits"); - - struct named_tensor { - const char* name; - struct ggml_tensor* tensor; - }; + auto* image_pe = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, D, H, H, 1); + ggml_set_name(image_pe, "prop_pe"); + ggml_set_input(image_pe); + auto* dense_emb = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, D, H, H, 1); + ggml_set_name(dense_emb, "prop_dense"); + ggml_set_input(dense_emb); - const std::vector direct_tensors = { - {"text_features", text_features}, - {"text_valid_mask", text_valid_mask}, - {"fenc_img_input", img_feats}, - {"fenc_pos_embed", img_pe}, - {"fenc_prompt", text_features}, - {"img_pe_72", img_pe_72_cmp}, - {"ddec_query_embed", model.ddec.query_embed}, - {"ddec_ref_pts_raw", model.tensors.at("ddec.reference_points.weight")}, - {"ddec_presence_token", model.ddec.presence_token}, - {"ddec_pred_boxes", ddec_out.pred_boxes}, - {"ddec_presence_logit", ddec_out.presence_score}, - }; + const int H0 = H * 4, H1 = H * 2; + auto* trk_s0 = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, D, H0, H0, 1); + ggml_set_name(trk_s0, "prop_trk_s0"); + ggml_set_input(trk_s0); + auto* trk_s1 = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, D, H1, H1, 1); + ggml_set_name(trk_s1, "prop_trk_s1"); + ggml_set_input(trk_s1); - std::vector named_outputs = { - "fenc_output", - "ddec_ref_boxes_init", - "ddec_query_sine_0", - "ddec_query_pos_0", - "ddec_rpb_mask_0", - "ddec_layer0_after_sa", - "ddec_layer0_after_text_ca", - "ddec_layer0_after_img_ca", - "ddec_layer0_full_out", - "ddec_layer0_presence", - "ddec_normed_output", - "scoring_prompt_mlp_out", - "scoring_pooled", - "scoring_proj_pooled", - "scoring_proj_hs", - "scoring_class_scores", - "seg_enc_after_ca", - "seg_enc_visual", - "seg_pixel_dec_stage0", - "seg_pixel_dec_stage1", - "seg_pixel_decoder_out", - "seg_instance_embed", - "seg_mask_embed", - "seg_mask_logits", - }; + auto dec = sam3_build_sam_dec_graph(ctx0, model, cond_spatial, image_pe, + sparse_in, dense_emb, + trk_s0, trk_s1, H); + ggml_set_output(dec.masks); + ggml_set_output(dec.iou_pred); + ggml_set_output(dec.obj_score); + ggml_set_output(dec.sam_token); + if (dec.mask_tokens) ggml_set_output(dec.mask_tokens); - for (int i = 0; i < hp.fenc_layers; ++i) { - char name[64]; - snprintf(name, sizeof(name), "fenc_layer%d_out", i); - named_outputs.emplace_back(name); - } - for (int i = 0; i < hp.ddec_layers; ++i) { - char out_name[64]; - char box_name[64]; - snprintf(out_name, sizeof(out_name), "ddec_layer%d_out", i); - snprintf(box_name, sizeof(box_name), "ddec_layer%d_refboxes", i); - named_outputs.emplace_back(out_name); - named_outputs.emplace_back(box_name); - } - - ggml_set_output(text_features); - ggml_set_output(text_valid_mask); - ggml_set_output(img_feats); - ggml_set_output(img_pe); - ggml_set_output(img_pe_72_cmp); - for (const auto& name : named_outputs) { - auto* t = ggml_get_tensor(ctx0, name.c_str()); - if (t) { - ggml_set_output(t); - } - } - ggml_set_output(mask_logits); - ggml_set_output(ddec_out.class_scores); - ggml_set_output(ddec_out.pred_boxes); - ggml_set_output(ddec_out.presence_score); - - struct ggml_cgraph* graph = ggml_new_graph_custom(ctx0, 65536, false); - ggml_build_forward_expand(graph, ddec_out.class_scores); - ggml_build_forward_expand(graph, ddec_out.pred_boxes); - ggml_build_forward_expand(graph, ddec_out.presence_score); - ggml_build_forward_expand(graph, mask_logits); - ggml_build_forward_expand(graph, text_features); - ggml_build_forward_expand(graph, text_valid_mask); - ggml_build_forward_expand(graph, img_feats); - ggml_build_forward_expand(graph, img_pe); - ggml_build_forward_expand(graph, img_pe_72_cmp); - for (const auto& name : named_outputs) { - auto* t = ggml_get_tensor(ctx0, name.c_str()); - if (t) { - ggml_build_forward_expand(graph, t); - } - } + auto* graph = ggml_new_graph_custom(ctx0, 32768, false); + ggml_build_forward_expand(graph, dec.masks); + ggml_build_forward_expand(graph, dec.iou_pred); + ggml_build_forward_expand(graph, dec.obj_score); + ggml_build_forward_expand(graph, dec.sam_token); + if (dec.mask_tokens) ggml_build_forward_expand(graph, dec.mask_tokens); auto* galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); if (!ggml_gallocr_reserve(galloc, graph) || !ggml_gallocr_alloc_graph(galloc, graph)) { - fprintf(stderr, "%s: failed to allocate phase 5 graph\n", __func__); ggml_gallocr_free(galloc); ggml_free(ctx0); - return false; + return output; } - ggml_backend_tensor_set(text_features, text_features_data.data(), 0, text_features_data.size() * sizeof(float)); - ggml_backend_tensor_set(neck0, neck_det_0.data(), 0, neck_det_0.size() * sizeof(float)); - ggml_backend_tensor_set(neck1, neck_det_1.data(), 0, neck_det_1.size() * sizeof(float)); - ggml_backend_tensor_set(img_feats, fenc_img_input_data.data(), 0, fenc_img_input_data.size() * sizeof(float)); - ggml_backend_tensor_set(img_pe, fenc_pos_embed_data.data(), 0, fenc_pos_embed_data.size() * sizeof(float)); - ggml_backend_tensor_set(img_pe_72_cmp, img_pe_72_data.data(), 0, img_pe_72_data.size() * sizeof(float)); - - { - std::vector dim_t_data(64); - for (int i = 0; i < 64; ++i) { - dim_t_data[i] = 2.0f * (float)M_PI / powf(10000.0f, 2.0f * (float)i / 128.0f); - } - ggml_backend_tensor_set(sine_dim_t, dim_t_data.data(), 0, 64 * sizeof(float)); - } + // Upload prompt data + ggml_backend_tensor_set(prompt_t, pd.prompt.data(), 0, pd.prompt.size() * sizeof(float)); + ggml_backend_tensor_set(prompt_pos_t, pd.prompt_pos.data(), 0, pd.prompt_pos.size() * sizeof(float)); + ggml_backend_tensor_set(rope_q_t, rope_q_reord.data(), 0, rope_q_reord.size() * sizeof(float)); + if (rope_k_t && !rope_k_data.empty()) + ggml_backend_tensor_set(rope_k_t, rope_k_data.data(), 0, rope_k_data.size() * sizeof(float)); - { - std::vector coords(H); - for (int i = 0; i < H; ++i) { - coords[i] = (float)i / (float)H; - } - ggml_backend_tensor_set(rpb_coords, coords.data(), 0, H * sizeof(float)); + // Set default obj_score when pred_obj_scores=False (older SAM2 models) + if (!model.sam_dec.obj_score_token) { + auto* t = ggml_graph_get_tensor(graph, "sam_dec_obj_score"); + if (t) { float v = 10.0f; ggml_backend_tensor_set(t, &v, 0, sizeof(float)); } } - auto* qpos_pres = ggml_get_tensor(ctx0, "ddec_query_pos_pres"); - if (qpos_pres) { - std::vector zeros(D, 0.0f); - ggml_backend_tensor_set(qpos_pres, zeros.data(), 0, D * sizeof(float)); - } + // Upload src_pos (sinusoidal PE 256-dim) + ggml_backend_tensor_set(src_pos_t, tracker.cached_sinpe_256.data(), 0, + tracker.cached_sinpe_256.size() * sizeof(float)); - auto* rpb_pz = ggml_get_tensor(ctx0, "rpb_pres_zeros"); - if (rpb_pz) { - int n = (int)(rpb_pz->ne[0] * rpb_pz->ne[1] * rpb_pz->ne[2] * rpb_pz->ne[3]); - std::vector zeros(n, 0.0f); - ggml_backend_tensor_set(rpb_pz, zeros.data(), 0, n * sizeof(float)); - } + // Upload not_a_point_embed, image_pe, dense_emb from state PE cache + sam3_populate_pe_cache(state, model); + ggml_backend_tensor_set(sparse_in, state.not_a_point_cache, 0, D * sizeof(float)); + ggml_backend_tensor_set(image_pe, state.dense_pe_cache.data(), 0, + state.dense_pe_cache.size() * sizeof(float)); + ggml_backend_tensor_set(dense_emb, state.dense_nomask_cache.data(), 0, + state.dense_nomask_cache.size() * sizeof(float)); + // Copy tracker features from state to fresh input tensors { - int n_valid = 0; - for (int i = 0; i < L; ++i) { - if (token_ids[i] != 0) { - ++n_valid; - } - } - if (n_valid == 0) { - n_valid = 1; - } - - const float scale = (float)L / (float)n_valid; - std::vector valid_mask(L); - std::vector attn_bias(L); - for (int i = 0; i < L; ++i) { - const bool is_valid = token_ids[i] != 0; - valid_mask[i] = is_valid ? scale : 0.0f; - attn_bias[i] = is_valid ? 0.0f : -1.0e9f; - } - ggml_backend_tensor_set(text_valid_mask, valid_mask.data(), 0, L * sizeof(float)); - ggml_backend_tensor_set(text_attn_bias, attn_bias.data(), 0, L * sizeof(float)); - } - - sam3_graph_compute(model.backend, graph, n_threads); - - bool ok = true; - for (const auto& item : direct_tensors) { - if (!sam3_dump_tensor_to_path(item.tensor, item.name, output_dir + "/" + item.name)) { - ok = false; - } - } - for (const auto& name : named_outputs) { - auto* t = ggml_get_tensor(ctx0, name.c_str()); - if (!t) { - fprintf(stderr, "%s: tensor '%s' not found in graph\n", __func__, name.c_str()); - ok = false; - continue; - } - if (!sam3_dump_tensor_to_path(t, name, output_dir + "/" + name)) { - ok = false; - } - } - - ggml_gallocr_free(galloc); - ggml_free(ctx0); - return ok; -} - -bool sam3_test_fenc_only(const sam3_model& model, - const std::string& ref_dir, - const std::string& output_dir, - int n_threads) { - const auto& hp = model.hparams; - const int D = hp.neck_dim; // 256 - const int H = hp.n_img_embd(); // 72 - const int N = H * H; // 5184 - - fprintf(stderr, "%s: D=%d H=%d N=%d fenc_layers=%d\n", - __func__, D, H, N, hp.fenc_layers); - - // ── Load reference tensors from Python dump ── - // The Python script saves batch-first [1, N, D] tensors. - // Memory layout: N blocks of D floats = same as ggml [D, N, 1]. - std::vector img_feat_data; - std::vector pos_data; - std::vector prompt_data; - std::vector attn_bias_data; - - if (!sam3_load_ref_f32_data(ref_dir + "/fenc_input_tgt", img_feat_data, D * N)) { - fprintf(stderr, "%s: failed to load fenc_input_tgt\n", __func__); - return false; - } - if (!sam3_load_ref_f32_data(ref_dir + "/fenc_input_pos", pos_data, D * N)) { - fprintf(stderr, "%s: failed to load fenc_input_pos\n", __func__); - return false; - } - if (!sam3_load_ref_f32_data(ref_dir + "/fenc_input_prompt", prompt_data)) { - fprintf(stderr, "%s: failed to load fenc_input_prompt\n", __func__); - return false; - } + std::vector c2(D * N); + ggml_backend_tensor_get(state.neck_trk[2], c2.data(), 0, D * N * sizeof(float)); + ggml_backend_tensor_set(curr, c2.data(), 0, D * N * sizeof(float)); - // Determine prompt length from loaded data - const int T = (int)prompt_data.size() / D; - fprintf(stderr, "%s: prompt tokens T=%d\n", __func__, T); + std::vector s0(D * H0 * H0); + ggml_backend_tensor_get(state.neck_trk[0], s0.data(), 0, D * H0 * H0 * sizeof(float)); + ggml_backend_tensor_set(trk_s0, s0.data(), 0, D * H0 * H0 * sizeof(float)); - // Attn bias: [1, T] float (0.0 valid, -1e9 padding) - // If not present, assume all valid (no bias) - bool have_attn_bias = sam3_load_ref_f32_data(ref_dir + "/fenc_attn_bias", attn_bias_data, T); - if (!have_attn_bias) { - fprintf(stderr, "%s: no fenc_attn_bias found, assuming all tokens valid\n", __func__); - attn_bias_data.assign(T, 0.0f); + std::vector s1(D * H1 * H1); + ggml_backend_tensor_get(state.neck_trk[1], s1.data(), 0, D * H1 * H1 * sizeof(float)); + ggml_backend_tensor_set(trk_s1, s1.data(), 0, D * H1 * H1 * sizeof(float)); } - // ── Build fenc-only graph ── - const size_t buf_size = ggml_tensor_overhead() * 16384 + ggml_graph_overhead() * 2; - struct ggml_init_params gparams = { - /*.mem_size =*/buf_size, - /*.mem_buffer =*/nullptr, - /*.no_alloc =*/true, - }; - struct ggml_context* ctx0 = ggml_init(gparams); - if (!ctx0) { - fprintf(stderr, "%s: failed to init compute context\n", __func__); - return false; + if (!sam3_graph_compute(model.backend, graph, 4)) { + ggml_gallocr_free(galloc); + ggml_free(ctx0); + return output; } - // Create input tensors (ggml layout: [D, N, 1]) - auto* img_feats = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, D, N, 1); - ggml_set_name(img_feats, "fenc_img_input"); - ggml_set_input(img_feats); - - auto* pos_enc = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, D, N, 1); - ggml_set_name(pos_enc, "fenc_pos_input"); - ggml_set_input(pos_enc); - - auto* prompt_tokens = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, D, T, 1); - ggml_set_name(prompt_tokens, "fenc_prompt_input"); - ggml_set_input(prompt_tokens); - - auto* text_attn_bias = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, T, 1, 1); - ggml_set_name(text_attn_bias, "fenc_attn_bias"); - ggml_set_input(text_attn_bias); - - // Build fusion encoder graph - auto* conditioned = sam3_build_fenc_graph(ctx0, model, img_feats, prompt_tokens, - pos_enc, text_attn_bias); - // sam3_build_fenc_graph names the last layer fenc_layer5_out. ggml_set_name - // will overwrite it, so create a cont copy so both names resolve. - auto* fenc_last = ggml_cont(ctx0, conditioned); - ggml_set_name(fenc_last, "fenc_layer5_out"); - ggml_set_name(conditioned, "fenc_output"); + const int mhw = H * 4; + const int num_mask_tokens = hp.sam_n_multimask + 1; // 4 - // Mark all per-layer outputs as graph outputs for extraction - std::vector output_names = {"fenc_output"}; - for (int i = 0; i < hp.fenc_layers; ++i) { - char name[64]; - snprintf(name, sizeof(name), "fenc_layer%d_out", i); - output_names.emplace_back(name); - } + // SAM2/EdgeTAM with multimask_output_in_sam: read all masks, select best by IoU. + // SAM3 / SAM2 without multimask: use first mask only. + bool use_multimask = (hp.is_sam2() || hp.is_edgetam()) && hp.multimask_output_in_sam; - ggml_set_output(img_feats); - ggml_set_output(pos_enc); - ggml_set_output(prompt_tokens); - for (const auto& name : output_names) { - auto* t = ggml_get_tensor(ctx0, name.c_str()); - if (t) { - ggml_set_output(t); - } - } + if (use_multimask) { + // Read all 4 IoU predictions + std::vector all_ious(num_mask_tokens); + ggml_backend_tensor_get(dec.iou_pred, all_ious.data(), 0, num_mask_tokens * sizeof(float)); - // Build and allocate graph - struct ggml_cgraph* graph = ggml_new_graph_custom(ctx0, 16384, false); - ggml_build_forward_expand(graph, conditioned); - ggml_build_forward_expand(graph, fenc_last); - for (const auto& name : output_names) { - auto* t = ggml_get_tensor(ctx0, name.c_str()); - if (t) { - ggml_build_forward_expand(graph, t); + // Multimask uses mask tokens 1-3 (skip token 0 which is the single-mask output) + int best_idx = 1; + float best_iou = all_ious[1]; + for (int m = 2; m < num_mask_tokens; ++m) { + if (all_ious[m] > best_iou) { best_iou = all_ious[m]; best_idx = m; } } - } - auto* galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); - if (!ggml_gallocr_reserve(galloc, graph) || !ggml_gallocr_alloc_graph(galloc, graph)) { - fprintf(stderr, "%s: failed to allocate fenc graph\n", __func__); - ggml_gallocr_free(galloc); - ggml_free(ctx0); - return false; - } + output.n_masks = 1; + output.mask_h = mhw; + output.mask_w = mhw; + output.mask_logits.resize(mhw * mhw); + // Read best mask (offset by best_idx * mhw * mhw) + ggml_backend_tensor_get(dec.masks, output.mask_logits.data(), + best_idx * mhw * mhw * sizeof(float), mhw * mhw * sizeof(float)); + output.iou_scores.resize(1); + output.iou_scores[0] = best_iou; + ggml_backend_tensor_get(dec.obj_score, &output.obj_score, 0, sizeof(float)); - // ── Set input data ── - ggml_backend_tensor_set(img_feats, img_feat_data.data(), 0, - img_feat_data.size() * sizeof(float)); - ggml_backend_tensor_set(pos_enc, pos_data.data(), 0, - pos_data.size() * sizeof(float)); - ggml_backend_tensor_set(prompt_tokens, prompt_data.data(), 0, - prompt_data.size() * sizeof(float)); - ggml_backend_tensor_set(text_attn_bias, attn_bias_data.data(), 0, - attn_bias_data.size() * sizeof(float)); - - // ── Compute ── - fprintf(stderr, "%s: computing fenc graph (%d nodes)...\n", __func__, ggml_graph_n_nodes(graph)); - sam3_graph_compute(model.backend, graph, n_threads); - - // ── Dump outputs ── - bool ok = true; - for (const auto& name : output_names) { - auto* t = ggml_get_tensor(ctx0, name.c_str()); - if (!t) { - fprintf(stderr, "%s: tensor '%s' not found\n", __func__, name.c_str()); - ok = false; - continue; - } - if (!sam3_dump_tensor_to_path(t, name, output_dir + "/" + name)) { - ok = false; + // Object pointer token: use best multimask token if use_multimask_token_for_obj_ptr + output.sam_token.resize(D); + if (hp.use_multimask_token_for_obj_ptr && dec.mask_tokens) { + ggml_backend_tensor_get(dec.mask_tokens, output.sam_token.data(), + best_idx * D * sizeof(float), D * sizeof(float)); + } else { + ggml_backend_tensor_get(dec.sam_token, output.sam_token.data(), 0, D * sizeof(float)); } + } else { + output.n_masks = 1; + output.mask_h = mhw; + output.mask_w = mhw; + output.mask_logits.resize(mhw * mhw); + ggml_backend_tensor_get(dec.masks, output.mask_logits.data(), 0, mhw * mhw * sizeof(float)); + output.iou_scores.resize(1); + ggml_backend_tensor_get(dec.iou_pred, output.iou_scores.data(), 0, sizeof(float)); + ggml_backend_tensor_get(dec.obj_score, &output.obj_score, 0, sizeof(float)); + output.sam_token.resize(D); + ggml_backend_tensor_get(dec.sam_token, output.sam_token.data(), 0, D * sizeof(float)); } - // Also dump inputs for verification - sam3_dump_tensor_to_path(img_feats, "fenc_img_input", output_dir + "/fenc_img_input"); - sam3_dump_tensor_to_path(pos_enc, "fenc_pos_input", output_dir + "/fenc_pos_input"); - sam3_dump_tensor_to_path(prompt_tokens, "fenc_prompt_input", output_dir + "/fenc_prompt_input"); - ggml_gallocr_free(galloc); ggml_free(ctx0); - return ok; -} - -static bool sam3_test_dump_phase6_impl(const sam3_model& model, - struct ggml_tensor* state_neck0, - struct ggml_tensor* state_neck1, - struct ggml_tensor* state_neck2, - const std::vector* neck0_data, - const std::vector* neck1_data, - const std::vector* neck2_data, - const sam3_pvs_params& params, - const std::string& output_dir, - int n_threads) { - const auto& hp = model.hparams; - const int D = hp.sam_embed_dim; - const int H = hp.n_img_embd(); - const int H1 = H * 2; - const int H0 = H * 4; + return output; +} - const size_t buf_size = ggml_tensor_overhead() * 32768 + ggml_graph_overhead() * 2; - struct ggml_init_params gparams = { - /*.mem_size =*/buf_size, - /*.mem_buffer =*/nullptr, - /*.no_alloc =*/true, - }; - struct ggml_context* ctx0 = ggml_init(gparams); - if (!ctx0) { - fprintf(stderr, "%s: failed to init compute context\n", __func__); - return false; +static std::vector> sam3_match_detections( + const std::vector& masklets, const std::vector& dets, + const std::vector& prop_masks, float iou_threshold) { + std::vector> matches; + if (masklets.empty() || dets.empty()) return matches; + int n_m = (int)masklets.size(), n_d = (int)dets.size(); + std::vector dm(n_d, false); + for (int i = 0; i < n_m; ++i) { + if (i >= (int)prop_masks.size() || prop_masks[i].data.empty()) continue; + int bj = -1; + float bi = iou_threshold; + for (int j = 0; j < n_d; ++j) { + if (dm[j] || dets[j].mask.data.empty()) continue; + int w = prop_masks[i].width, h = prop_masks[i].height; + if (w != dets[j].mask.width || h != dets[j].mask.height) continue; + float iou = sam3_mask_iou(prop_masks[i].data.data(), dets[j].mask.data.data(), w * h); + if (iou > bi) { + bi = iou; + bj = j; + } + } + if (bj >= 0) { + matches.push_back({i, bj}); + dm[bj] = true; + } } + return matches; +} - struct ggml_tensor* neck0 = state_neck0; - struct ggml_tensor* neck1 = state_neck1; - struct ggml_tensor* neck2 = state_neck2; - if (neck0_data && neck1_data && neck2_data) { - neck0 = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, D, H0, H0, 1); - ggml_set_name(neck0, "sam_dec_feat_s0_input"); - ggml_set_input(neck0); - - neck1 = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, D, H1, H1, 1); - ggml_set_name(neck1, "sam_dec_feat_s1_input"); - ggml_set_input(neck1); - - neck2 = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, D, H, H, 1); - ggml_set_name(neck2, "sam_dec_feat_s2_input"); - ggml_set_input(neck2); +static void sam3_update_tracker(sam3_tracker& tracker, int frame_idx) { + for (auto it = tracker.pending.begin(); it != tracker.pending.end();) { + int age = frame_idx - it->first_frame; + if (age >= tracker.params.hotstart_delay && it->mds_sum > 0) { + it->confirmed = true; + tracker.masklets.push_back(std::move(*it)); + it = tracker.pending.erase(it); + } else if (age >= tracker.params.hotstart_delay) { + it = tracker.pending.erase(it); + } else + ++it; } - - if (!neck0 || !neck1 || !neck2) { - fprintf(stderr, "%s: missing tracker feature inputs\n", __func__); - ggml_free(ctx0); - return false; + for (auto it = tracker.masklets.begin(); it != tracker.masklets.end();) { + if (frame_idx - it->last_seen > tracker.params.max_keep_alive) { + tracker.mem_banks.erase(it->instance_id); + tracker.ptr_banks.erase(it->instance_id); + it = tracker.masklets.erase(it); + } else + ++it; } +} - auto pe_out = sam3_build_sam_pe(ctx0, params, D, H); - - auto* no_mem = ggml_reshape_4d(ctx0, model.tensors.at("no_mem_embed"), D, 1, 1, 1); - auto* image_feats = ggml_add(ctx0, neck2, no_mem); - ggml_set_name(image_feats, "sam_dec_image_feats"); - - auto dec_out = sam3_build_sam_dec_graph(ctx0, model, image_feats, - pe_out.image_pe, pe_out.sparse, pe_out.dense, - neck0, neck1); +static bool sam3_encode_memory( + sam3_tracker& tracker, sam3_state& state, const sam3_model& model, + int inst_id, const float* mask_logits, int mask_h, int mask_w, + int frame_idx, bool is_cond, float obj_score) { + const auto& hp = model.hparams; + const int D = hp.neck_dim, MD = hp.mem_out_dim; + const int H = sam3_eff_feat_size(state, hp); + const int HIGH_RES = sam3_eff_img_size(state, hp); + const int INTERPOL = H * 16; - struct named_tensor { - const char* name; - struct ggml_tensor* tensor; - }; + // Mask preprocessing: mask_logits → HIGH_RES → sigmoid → scale/bias → INTERPOL + auto m_hires = sam3_bilinear_interpolate(mask_logits, mask_w, mask_h, HIGH_RES, HIGH_RES); + const float sig_scale = hp.sigmoid_scale(), sig_bias = hp.sigmoid_bias(); + for (auto& v : m_hires) { float s = 1.0f / (1.0f + expf(-v)); v = s * sig_scale + sig_bias; } + auto m_interp = sam3_bilinear_interpolate(m_hires.data(), HIGH_RES, HIGH_RES, INTERPOL, INTERPOL); - const std::vector direct_tensors = { - {"sam_pe_sparse", pe_out.sparse}, - {"sam_pe_dense", pe_out.dense}, - {"sam_pe_image_pe", pe_out.image_pe}, - {"sam_dec_image_feats", image_feats}, - {"sam_dec_sam_token", dec_out.sam_token}, - }; + const size_t bs = ggml_tensor_overhead() * 16384 + ggml_graph_overhead(); + struct ggml_init_params gp = {bs, nullptr, true}; + auto* ctx0 = ggml_init(gp); + if (!ctx0) return false; - std::vector named_outputs = { - "sam_dec_tokens_initial", - "sam_dec_block0_queries", - "sam_dec_block0_keys", - "sam_dec_block1_queries", - "sam_dec_block1_keys", - "sam_dec_final_queries", - "sam_dec_feat_s1_proj", - "sam_dec_feat_s0_proj", - "sam_dec_upscaled", - "sam_dec_mask_tokens", - "sam_dec_masks", - "sam_dec_iou", - "sam_dec_obj_score", - }; + auto* mask_in = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, INTERPOL, INTERPOL, 1, 1); + ggml_set_name(mask_in, "mem_mask"); + ggml_set_input(mask_in); - ggml_set_output(dec_out.masks); - ggml_set_output(dec_out.iou_pred); - ggml_set_output(dec_out.obj_score); - ggml_set_output(dec_out.sam_token); - ggml_set_output(dec_out.mask_tokens); - for (const auto& item : direct_tensors) { - ggml_set_output(item.tensor); - } - for (const auto& name : named_outputs) { - auto* t = ggml_get_tensor(ctx0, name.c_str()); - if (!t) { - fprintf(stderr, "%s: tensor '%s' not found in graph\n", __func__, name.c_str()); - ggml_free(ctx0); - return false; - } - ggml_set_output(t); + // Learned mask downsampler: 4 stages conv+LN+GELU + final 1×1 + // Conv output is in ggml "conv" layout. permute(1,2,0,3) converts to "internal" layout + // where ne[0]=channel for LN2d. permute(2,0,1,3) converts back to "conv" layout. + auto* ds = mask_in; + for (int s = 0; s < 4; ++s) { + int out_ch = (int)model.mem_enc.ds_conv_w[s]->ne[3]; + ds = ggml_conv_2d(ctx0, model.mem_enc.ds_conv_w[s], ds, 2, 2, 1, 1, 1, 1); + ds = ggml_add(ctx0, ds, ggml_reshape_4d(ctx0, model.mem_enc.ds_conv_b[s], 1, 1, out_ch, 1)); + ds = ggml_cont(ctx0, ggml_permute(ctx0, ds, 1, 2, 0, 3)); // conv→internal (ne[0]=channel) + ds = sam3_layer_norm_2d(ctx0, ds, model.mem_enc.ds_norm_w[s], model.mem_enc.ds_norm_b[s]); + ds = ggml_gelu(ctx0, ds); + ds = ggml_cont(ctx0, ggml_permute(ctx0, ds, 2, 0, 1, 3)); // internal→conv } + ds = ggml_conv_2d(ctx0, model.mem_enc.ds_conv_w[4], ds, 1, 1, 0, 0, 1, 1); + ds = ggml_add(ctx0, ds, ggml_reshape_4d(ctx0, model.mem_enc.ds_conv_b[4], 1, 1, D, 1)); + ds = ggml_cont(ctx0, ggml_permute(ctx0, ds, 1, 2, 0, 3)); // conv→internal [D, ...] - struct ggml_cgraph* graph = ggml_new_graph_custom(ctx0, 32768, false); - ggml_build_forward_expand(graph, dec_out.masks); - ggml_build_forward_expand(graph, dec_out.iou_pred); - ggml_build_forward_expand(graph, dec_out.obj_score); - ggml_build_forward_expand(graph, dec_out.sam_token); - ggml_build_forward_expand(graph, dec_out.mask_tokens); - for (const auto& item : direct_tensors) { - ggml_build_forward_expand(graph, item.tensor); - } - for (const auto& name : named_outputs) { - ggml_build_forward_expand(graph, ggml_get_tensor(ctx0, name.c_str())); - } + // Pixel projection — use fresh input tensor to avoid pulling in the + // entire ViT+neck recomputation from state.neck_trk[2]'s dependency tree + auto* pix_in_raw = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, D, H, H, 1); + ggml_set_name(pix_in_raw, "mem_pix_feat"); + ggml_set_input(pix_in_raw); + auto* pix_in = ggml_cont(ctx0, ggml_permute(ctx0, pix_in_raw, 2, 0, 1, 3)); + auto* pix = ggml_conv_2d(ctx0, model.mem_enc.pix_proj_w, pix_in, 1, 1, 0, 0, 1, 1); + pix = ggml_add(ctx0, pix, ggml_reshape_4d(ctx0, model.mem_enc.pix_proj_b, 1, 1, D, 1)); + pix = ggml_cont(ctx0, ggml_permute(ctx0, pix, 1, 2, 0, 3)); - auto* galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); - if (!ggml_gallocr_reserve(galloc, graph) || !ggml_gallocr_alloc_graph(galloc, graph)) { - fprintf(stderr, "%s: failed to allocate phase 6 graph\n", __func__); - ggml_gallocr_free(galloc); + // Fusion: ADD (not multiply) + auto* fused = ggml_add(ctx0, pix, ds); + for (int i = 0; i < 2; ++i) + fused = sam3_cxblock_forward(ctx0, fused, + model.mem_enc.fuser_dw_w[i], model.mem_enc.fuser_dw_b[i], + model.mem_enc.fuser_norm_w[i], model.mem_enc.fuser_norm_b[i], + model.mem_enc.fuser_fc1_w[i], model.mem_enc.fuser_fc1_b[i], + model.mem_enc.fuser_fc2_w[i], model.mem_enc.fuser_fc2_b[i], + model.mem_enc.fuser_gamma[i]); + auto* fused_out = ggml_cont(ctx0, ggml_permute(ctx0, fused, 2, 0, 1, 3)); + auto* mo = ggml_conv_2d(ctx0, model.mem_enc.out_proj_w, fused_out, 1, 1, 0, 0, 1, 1); + mo = ggml_add(ctx0, mo, ggml_reshape_4d(ctx0, model.mem_enc.out_proj_b, 1, 1, MD, 1)); + mo = ggml_cont(ctx0, ggml_permute(ctx0, mo, 1, 2, 0, 3)); + ggml_set_name(mo, "mem_out"); + ggml_set_output(mo); + + auto* g = ggml_new_graph_custom(ctx0, 16384, false); + ggml_build_forward_expand(g, mo); + auto* ga = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); + if (!ggml_gallocr_reserve(ga, g) || !ggml_gallocr_alloc_graph(ga, g)) { + ggml_gallocr_free(ga); ggml_free(ctx0); return false; } - - if (neck0_data && neck1_data && neck2_data) { - ggml_backend_tensor_set(neck0, neck0_data->data(), 0, neck0_data->size() * sizeof(float)); - ggml_backend_tensor_set(neck1, neck1_data->data(), 0, neck1_data->size() * sizeof(float)); - ggml_backend_tensor_set(neck2, neck2_data->data(), 0, neck2_data->size() * sizeof(float)); + ggml_backend_tensor_set(mask_in, m_interp.data(), 0, m_interp.size() * sizeof(float)); + // Copy pixel features from state tensor to fresh input + { + std::vector pix_data(D * H * H); + ggml_backend_tensor_get(state.neck_trk[2], pix_data.data(), 0, D * H * H * sizeof(float)); + ggml_backend_tensor_set(pix_in_raw, pix_data.data(), 0, D * H * H * sizeof(float)); } - - sam3_state pe_cache_state = {}; - sam3_populate_pe_cache(pe_cache_state, model); - - std::vector all_coords; - std::vector all_labels; - sam3_collect_pvs_prompt_tokens(params, all_coords, all_labels); - if ((int)all_labels.size() != pe_out.n_tokens) { - fprintf(stderr, "%s: prompt token count mismatch (%zu vs %d)\n", - __func__, all_labels.size(), pe_out.n_tokens); - ggml_gallocr_free(galloc); + if (!sam3_graph_compute(model.backend, g, 4)) { + ggml_gallocr_free(ga); ggml_free(ctx0); return false; } - { - const int num_pos_feats = D / 2; - std::vector sparse_data((size_t)pe_out.n_tokens * D, 0.0f); - for (int p = 0; p < pe_out.n_tokens; ++p) { - float px = all_coords[p * 2 + 0] + 0.5f; - float py = all_coords[p * 2 + 1] + 0.5f; - float x_norm = px / (float)hp.img_size; - float y_norm = py / (float)hp.img_size; - float pe_vec[256]; - sam3_pe_encode_coord(pe_vec, x_norm, y_norm, - pe_cache_state.pe_gauss_cache.data(), num_pos_feats); + std::vector md(MD * H * H); + ggml_backend_tensor_get(mo, md.data(), 0, md.size() * sizeof(float)); - const int label = all_labels[p]; - if (label == -1) { - for (int d = 0; d < D; ++d) { - sparse_data[(size_t)p * D + d] = pe_cache_state.not_a_point_cache[d]; - } - } else { - for (int d = 0; d < D; ++d) { - sparse_data[(size_t)p * D + d] = pe_vec[d] + pe_cache_state.point_emb_cache[label][d]; - } - } + // Apply no_obj_embed_spatial if occluded (SAM2.1 only — EdgeTAM does not have this) + if (obj_score <= 0.0f && model.no_obj_embed_spatial) { + std::vector no_obj_emb(MD); + auto* noe = model.no_obj_embed_spatial; + if (noe->type == GGML_TYPE_F16) { + std::vector tmp(MD); + ggml_backend_tensor_get(noe, tmp.data(), 0, MD * sizeof(ggml_fp16_t)); + ggml_fp16_to_fp32_row(tmp.data(), no_obj_emb.data(), MD); + } else { + ggml_backend_tensor_get(noe, no_obj_emb.data(), 0, MD * sizeof(float)); } + for (int i = 0; i < MD * H * H; ++i) + md[i] += no_obj_emb[i % MD]; + } - ggml_backend_tensor_set(pe_out.sparse, sparse_data.data(), 0, - sparse_data.size() * sizeof(float)); - ggml_backend_tensor_set(pe_out.image_pe, pe_cache_state.dense_pe_cache.data(), 0, - pe_cache_state.dense_pe_cache.size() * sizeof(float)); - ggml_backend_tensor_set(pe_out.dense, pe_cache_state.dense_nomask_cache.data(), 0, - pe_cache_state.dense_nomask_cache.size() * sizeof(float)); + if (!tracker.ctx) { + struct ggml_init_params tp = {ggml_tensor_overhead() * 4096, nullptr, true}; + tracker.ctx = ggml_init(tp); } - sam3_graph_compute(model.backend, graph, n_threads); + // ── EdgeTAM perceiver: compress memory features before storage ─────── + if (hp.has_perceiver) { + // Compute sinusoidal PE for memory features (needed as perceiver input) + sam3_ensure_tracker_pe_caches(tracker, hp, H); + const auto& mem_pos = tracker.cached_sinpe_64; - bool ok = true; - for (const auto& item : direct_tensors) { - if (!sam3_dump_tensor_to_path(item.tensor, item.name, output_dir + "/" + item.name)) { - ok = false; - } - } - for (const auto& name : named_outputs) { - auto* t = ggml_get_tensor(ctx0, name.c_str()); - if (!t) { - fprintf(stderr, "%s: tensor '%s' not found after compute\n", __func__, name.c_str()); - ok = false; - continue; - } - if (!sam3_dump_tensor_to_path(t, name, output_dir + "/" + name)) { - ok = false; + std::vector perc_latents, perc_pos; + if (!edgetam_perceiver_forward(model, md, mem_pos, H, H, + perc_latents, perc_pos)) { + fprintf(stderr, "%s: perceiver forward failed\n", __func__); + ggml_gallocr_free(ga); + ggml_free(ctx0); + return false; } - } - ggml_gallocr_free(galloc); - ggml_free(ctx0); - return ok; -} + // Store perceiver output: [MD=64, N_perc] (N_1d + N_2d latents) + const int N_perc = hp.perceiver_n_latents_1d + hp.perceiver_n_latents_2d; + auto* st = ggml_new_tensor_2d(tracker.ctx, GGML_TYPE_F32, MD, N_perc); + auto* sb = ggml_backend_alloc_buffer(model.backend, MD * N_perc * sizeof(float)); + struct ggml_tallocr ta_perc = ggml_tallocr_new(sb); + ggml_tallocr_alloc(&ta_perc, st); + tracker.owned_buffers.push_back(sb); + ggml_backend_tensor_set(st, perc_latents.data(), 0, MD * N_perc * sizeof(float)); -bool sam3_test_dump_phase6(const sam3_model& model, - const sam3_state& state, - const sam3_pvs_params& params, - const std::string& output_dir, - int n_threads) { - return sam3_test_dump_phase6_impl(model, - state.neck_trk[0], state.neck_trk[1], state.neck_trk[2], - nullptr, nullptr, nullptr, - params, output_dir, n_threads); -} + // Store perceiver PE: [MD=64, 512] (first 256 zeros, last 256 sinusoidal) + auto* spe = ggml_new_tensor_2d(tracker.ctx, GGML_TYPE_F32, MD, N_perc); + auto* speb = ggml_backend_alloc_buffer(model.backend, MD * N_perc * sizeof(float)); + struct ggml_tallocr ta_perc2 = ggml_tallocr_new(speb); + ggml_tallocr_alloc(&ta_perc2, spe); + tracker.owned_buffers.push_back(speb); + ggml_backend_tensor_set(spe, perc_pos.data(), 0, MD * N_perc * sizeof(float)); -bool sam3_test_dump_phase6_from_ref_inputs(const sam3_model& model, - const std::string& prephase_ref_dir, - const sam3_pvs_params& params, - const std::string& output_dir, - int n_threads) { - const auto& hp = model.hparams; - const int D = hp.sam_embed_dim; - const int H = hp.n_img_embd(); - const int H1 = H * 2; - const int H0 = H * 4; - - std::vector neck_trk_0; - std::vector neck_trk_1; - std::vector neck_trk_2; - if (!sam3_load_ref_f32_data(prephase_ref_dir + "/neck_trk_0", neck_trk_0, D * H0 * H0) || - !sam3_load_ref_f32_data(prephase_ref_dir + "/neck_trk_1", neck_trk_1, D * H1 * H1) || - !sam3_load_ref_f32_data(prephase_ref_dir + "/neck_trk_2", neck_trk_2, D * H * H)) { - return false; + sam3_memory_slot slot; + slot.spatial_feats = st; + slot.spatial_pe = spe; + slot.frame_index = frame_idx; + slot.is_cond_frame = is_cond; + auto& bk = tracker.mem_banks[inst_id]; + bk.push_back(slot); + while ((int)bk.size() > hp.num_maskmem) { + bool removed = false; + for (auto it = bk.begin(); it != bk.end(); ++it) + if (!it->is_cond_frame) { + bk.erase(it); + removed = true; + break; + } + if (!removed) bk.erase(bk.begin() + 1); + } + ggml_gallocr_free(ga); + ggml_free(ctx0); + return true; } - neck_trk_0 = sam3_reorder_nchw_to_ggml_dwh(neck_trk_0, D, H0, H0); - neck_trk_1 = sam3_reorder_nchw_to_ggml_dwh(neck_trk_1, D, H1, H1); - neck_trk_2 = sam3_reorder_nchw_to_ggml_dwh(neck_trk_2, D, H, H); - - return sam3_test_dump_phase6_impl(model, - nullptr, nullptr, nullptr, - &neck_trk_0, &neck_trk_1, &neck_trk_2, - params, output_dir, n_threads); -} - -/***************************************************************************** -** Test: Geometry encoder dump -*****************************************************************************/ - -bool sam3_test_dump_geom_enc(const sam3_model& model, - const std::string& prephase_ref_dir, - const sam3_pcs_params& params, - const std::string& output_dir, - int n_threads) { - const auto& hp = model.hparams; - const int D = hp.neck_dim; // 256 - const int H = hp.n_img_embd(); // 72 + // ── Standard path (SAM2/SAM3): store raw [MD, H, H] features ──────── + // Store spatial features + auto* st = ggml_new_tensor_4d(tracker.ctx, GGML_TYPE_F32, MD, H, H, 1); + auto* sb = ggml_backend_alloc_buffer(model.backend, MD * H * H * sizeof(float)); + struct ggml_tallocr ta = ggml_tallocr_new(sb); + ggml_tallocr_alloc(&ta, st); + tracker.owned_buffers.push_back(sb); + ggml_backend_tensor_set(st, md.data(), 0, md.size() * sizeof(float)); - // Load backbone features from Phase 3 reference (NCHW format) - std::vector neck_det_2; - if (!sam3_load_ref_f32_data(prephase_ref_dir + "/neck_det_2", neck_det_2, D * H * H)) { - fprintf(stderr, "%s: failed to load neck_det_2\n", __func__); - return false; - } - // Reorder from NCHW to ggml [D, W, H] layout - auto neck_det_2_ggml = sam3_reorder_nchw_to_ggml_dwh(neck_det_2, D, H, H); + // Compute and store sinusoidal spatial PE + sam3_ensure_tracker_pe_caches(tracker, hp, H); + const auto& pe_data = tracker.cached_sinpe_64; + auto* spe = ggml_new_tensor_4d(tracker.ctx, GGML_TYPE_F32, MD, H, H, 1); + auto* speb = ggml_backend_alloc_buffer(model.backend, MD * H * H * sizeof(float)); + struct ggml_tallocr ta2 = ggml_tallocr_new(speb); + ggml_tallocr_alloc(&ta2, spe); + tracker.owned_buffers.push_back(speb); + ggml_backend_tensor_set(spe, pe_data.data(), 0, pe_data.size() * sizeof(float)); - // Compute sinusoidal PE for image features (matching Python PositionEmbeddingSine) - std::vector img_pe_nchw(D * H * H); - { - const float scale = 2.0f * (float)M_PI; - const float eps = 1e-6f; - const int num_pos_feats = D / 2; // 128 - - for (int row = 0; row < H; ++row) { - for (int col = 0; col < H; ++col) { - float y_embed = ((float)(row + 1) - 0.5f) / ((float)H + eps) * scale; - float x_embed = ((float)(col + 1) - 0.5f) / ((float)H + eps) * scale; - - for (int i = 0; i < num_pos_feats; ++i) { - int div_idx = 2 * (i / 2); - float dim_t = powf(10000.0f, (float)div_idx / (float)num_pos_feats); - - float px = x_embed / dim_t; - float py = y_embed / dim_t; - - float pe_y, pe_x; - if (i % 2 == 0) { - pe_y = sinf(py); - pe_x = sinf(px); - } else { - pe_y = cosf(py); - pe_x = cosf(px); - } - // Output: [1, D, H, W] NCHW — pos_y first half, pos_x second half - img_pe_nchw[i * H * H + row * H + col] = pe_y; - img_pe_nchw[(num_pos_feats + i) * H * H + row * H + col] = pe_x; - } + sam3_memory_slot slot; + slot.spatial_feats = st; + slot.spatial_pe = spe; + slot.frame_index = frame_idx; + slot.is_cond_frame = is_cond; + auto& bk = tracker.mem_banks[inst_id]; + bk.push_back(slot); + while ((int)bk.size() > hp.num_maskmem) { + bool removed = false; + for (auto it = bk.begin(); it != bk.end(); ++it) + if (!it->is_cond_frame) { + bk.erase(it); + removed = true; + break; } - } - } - auto img_pe_ggml = sam3_reorder_nchw_to_ggml_dwh(img_pe_nchw, D, H, H); - - // ── Build graph ───────────────────────────────────────────────────── - const size_t buf_size = ggml_tensor_overhead() * 8192 + ggml_graph_overhead() * 2; - struct ggml_init_params gparams = { - /*.mem_size =*/buf_size, - /*.mem_buffer =*/nullptr, - /*.no_alloc =*/true, - }; - struct ggml_context* ctx0 = ggml_init(gparams); - if (!ctx0) { - fprintf(stderr, "%s: failed to init context\n", __func__); - return false; + if (!removed) bk.erase(bk.begin() + 1); } + ggml_gallocr_free(ga); + ggml_free(ctx0); + return true; +} - // Image features as input tensors (from reference data) - auto* img_feats_t = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, D, H * H, 1); - ggml_set_name(img_feats_t, "img_feats_input"); - ggml_set_input(img_feats_t); - - auto* img_pe_t = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, D, H * H, 1); - ggml_set_name(img_pe_t, "img_pe_input"); - ggml_set_input(img_pe_t); - - // Build geometry encoder graph - auto geom_out = sam3_build_geom_enc_graph(ctx0, model, params, img_feats_t, img_pe_t); - - // Build and allocate graph - struct ggml_cgraph* graph = ggml_new_graph_custom(ctx0, 16384, false); - ggml_build_forward_expand(graph, geom_out.geo_feats); - - auto* galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); - if (!ggml_gallocr_reserve(galloc, graph) || !ggml_gallocr_alloc_graph(galloc, graph)) { - fprintf(stderr, "%s: failed to allocate graph\n", __func__); - ggml_gallocr_free(galloc); - ggml_free(ctx0); - return false; +static void sam3_store_obj_ptr( + sam3_tracker& tracker, const sam3_model& model, + int inst_id, const float* pd, int frame_idx) { + const int D = model.hparams.neck_dim; + if (!tracker.ctx) { + struct ggml_init_params tp = {ggml_tensor_overhead() * 4096, nullptr, true}; + tracker.ctx = ggml_init(tp); } + auto* pt = ggml_new_tensor_2d(tracker.ctx, GGML_TYPE_F32, D, 1); + auto* pb = ggml_backend_alloc_buffer(model.backend, D * sizeof(float)); + struct ggml_tallocr ta = ggml_tallocr_new(pb); + ggml_tallocr_alloc(&ta, pt); + tracker.owned_buffers.push_back(pb); + ggml_backend_tensor_set(pt, pd, 0, D * sizeof(float)); + auto& bk = tracker.ptr_banks[inst_id]; + bk.push_back({frame_idx, pt}); + while ((int)bk.size() > model.hparams.max_obj_ptrs) bk.erase(bk.begin()); +} - // Upload image features - ggml_backend_tensor_set(img_feats_t, neck_det_2_ggml.data(), 0, D * H * H * sizeof(float)); - ggml_backend_tensor_set(img_pe_t, img_pe_ggml.data(), 0, D * H * H * sizeof(float)); - - // Pre-compute geometry input on CPU and upload - auto geom_data = sam3_precompute_geom_input(model, params, neck_det_2_ggml.data(), H, H); - auto* gi = ggml_get_tensor(ctx0, "geom_post_final_proj"); - if (gi) { - ggml_backend_tensor_set(gi, geom_data.data(), 0, geom_data.size() * sizeof(float)); +sam3_tracker_ptr sam3_create_tracker(const sam3_model& model, + const sam3_video_params& params) { + if (model.hparams.is_sam2()) { + fprintf(stderr, "%s: ERROR: text-prompted tracker not available for SAM2 " + "(use sam3_create_visual_tracker instead)\n", __func__); + return nullptr; } + sam3_tracker_ptr tracker(new sam3_tracker()); + tracker->params = params; + fprintf(stderr, "%s: tracker created (hotstart=%d, max_keep_alive=%d)\n", + __func__, params.hotstart_delay, params.max_keep_alive); + return tracker; +} - // Compute - sam3_graph_compute(model.backend, graph, n_threads); - - // Dump outputs - mkdir(output_dir.c_str(), 0755); - - // Dump the pre-computed input (after final_proj + norm, before transformer) - { - auto* pre_proj = ggml_get_tensor(ctx0, "geom_post_final_proj"); - if (pre_proj) sam3_dump_tensor_to_path(pre_proj, "post_final_proj", output_dir + "/post_final_proj"); +sam3_result sam3_track_frame(sam3_tracker& tracker, sam3_state& state, + const sam3_model& model, const sam3_image& frame) { + if (model.hparams.visual_only) { + fprintf(stderr, "%s: ERROR: track_frame not available on visual-only model " + "(use sam3_propagate_frame instead)\n", __func__); + return sam3_result{}; } - // Dump final output - { - auto* out = ggml_get_tensor(ctx0, "geom_output"); - if (out) sam3_dump_tensor_to_path(out, "geom_output", output_dir + "/geom_output"); - } + sam3_result result; + const int D = model.hparams.neck_dim; + if (!sam3_encode_image(state, model, frame)) return result; + int fi = tracker.frame_index; + fprintf(stderr, "%s: frame %d (%zu active + %zu pending)\n", + __func__, fi, tracker.masklets.size(), tracker.pending.size()); - // Also dump the pre-computed input data (before ggml processing) - { - const int N_geo = geom_out.n_tokens; - std::string path = output_dir + "/geom_input_precomputed"; - { - std::ofstream f(path + ".bin", std::ios::binary); - f.write(reinterpret_cast(geom_data.data()), geom_data.size() * sizeof(float)); + std::map pm; + std::map po; + for (auto& ml : tracker.masklets) { + int id = ml.instance_id; + auto im = tracker.mem_banks.find(id); + if (im == tracker.mem_banks.end() || im->second.empty()) continue; + po[id] = sam3_propagate_single(tracker, state, model, ml, im->second, tracker.ptr_banks[id]); + if (po[id].mask_logits.empty()) continue; + auto rs = sam3_bilinear_interpolate(po[id].mask_logits.data(), + po[id].mask_w, po[id].mask_h, state.orig_width, state.orig_height); + pm[id].width = state.orig_width; + pm[id].height = state.orig_height; + pm[id].data.resize(state.orig_width * state.orig_height); + int fg = 0; + for (int p = 0; p < (int)rs.size(); ++p) { + bool f = rs[p] > 0.0f; + pm[id].data[p] = f ? 255 : 0; + if (f) fg++; } - { - std::ofstream f(path + ".shape"); - f << D << "," << N_geo; + ml.last_score = po[id].iou_scores[0]; + ml.last_seen = fi; + float cov = (float)fg / (state.orig_width * state.orig_height); + ml.mds_sum += (cov > 0.001f && po[id].obj_score > 0.0f) ? 1 : -1; + } + for (auto& ml : tracker.pending) { + int id = ml.instance_id; + auto im = tracker.mem_banks.find(id); + if (im == tracker.mem_banks.end() || im->second.empty()) continue; + auto p2 = sam3_propagate_single(tracker, state, model, ml, im->second, tracker.ptr_banks[id]); + if (!p2.mask_logits.empty()) { + ml.last_score = p2.iou_scores[0]; + ml.last_seen = fi; + auto r2 = sam3_bilinear_interpolate(p2.mask_logits.data(), + p2.mask_w, p2.mask_h, state.orig_width, state.orig_height); + int fg2 = 0; + for (auto v : r2) + if (v > 0.0f) fg2++; + float c2 = (float)fg2 / (state.orig_width * state.orig_height); + ml.mds_sum += (c2 > 0.001f && p2.obj_score > 0.0f) ? 1 : -1; + sam3_encode_memory(tracker, state, model, id, + p2.mask_logits.data(), p2.mask_h, p2.mask_w, fi, false, p2.obj_score); + std::vector op(D); + sam3_extract_obj_ptr_cpu(model, p2.sam_token.data(), p2.obj_score, op.data()); + sam3_store_obj_ptr(tracker, model, id, op.data(), fi); } - fprintf(stderr, "%s: dumped geom_input_precomputed [%d, %d]\n", __func__, D, N_geo); } - - ggml_gallocr_free(galloc); - ggml_free(ctx0); - return true; -} - -static bool sam3_test_dump_phase7_mem_slot_current( - const sam3_model& model, - const std::vector& neck2_data, - const std::vector& low_res_mask_logits, - float obj_score, - int slot_idx, - const std::string& output_dir, - std::vector& mem_out_data, - std::vector& mem_pe_data, - int n_threads) { - const auto& hp = model.hparams; - const int D = hp.neck_dim; - const int MD = hp.mem_out_dim; - const int H = hp.feat_size(); - const int MASK_HW = H * 4; - const int HIGH_RES = hp.img_size; - const int INTERPOL = H * 16; - - // Mask preprocessing: MASK_HW → HIGH_RES → sigmoid+scale+bias → INTERPOL - auto m_hires = sam3_bilinear_interpolate(low_res_mask_logits.data(), MASK_HW, MASK_HW, HIGH_RES, HIGH_RES); - const float sig_s = hp.sigmoid_scale(), sig_b = hp.sigmoid_bias(); - for (auto& v : m_hires) { float s = 1.0f / (1.0f + expf(-v)); v = s * sig_s + sig_b; } - auto m_interp = sam3_bilinear_interpolate(m_hires.data(), HIGH_RES, HIGH_RES, INTERPOL, INTERPOL); - - const size_t buf_size = ggml_tensor_overhead() * 16384 + ggml_graph_overhead(); - struct ggml_init_params gparams = {buf_size, nullptr, true}; - struct ggml_context* ctx0 = ggml_init(gparams); - if (!ctx0) { - fprintf(stderr, "%s: failed to init memory encoder context\n", __func__); - return false; + sam3_result nd; + if (!tracker.params.text_prompt.empty()) { + sam3_pcs_params pcs; + pcs.text_prompt = tracker.params.text_prompt; + pcs.score_threshold = tracker.params.score_threshold; + pcs.nms_threshold = tracker.params.nms_threshold; + nd = sam3_segment_pcs(state, model, pcs); } - - auto* neck2 = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, D, H, H, 1); - ggml_set_name(neck2, "phase7_curr_neck2"); - ggml_set_input(neck2); - - // Mask input at interpol resolution in WHCB layout for ggml_conv_2d - auto* mask_in = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, INTERPOL, INTERPOL, 1, 1); - ggml_set_name(mask_in, "phase7_mask_interp"); - ggml_set_input(mask_in); - - // ── Learned mask downsampler: 4 stages conv+LN+GELU, then 1×1 conv ── - // permute(1,2,0,3) converts conv→internal (ne[0]=channel for LN2d) - // permute(2,0,1,3) converts internal→conv - auto* ds = mask_in; - for (int s = 0; s < 4; ++s) { - int out_ch = (int)model.mem_enc.ds_conv_w[s]->ne[3]; - ds = ggml_conv_2d(ctx0, model.mem_enc.ds_conv_w[s], ds, 2, 2, 1, 1, 1, 1); - ds = ggml_add(ctx0, ds, ggml_reshape_4d(ctx0, model.mem_enc.ds_conv_b[s], 1, 1, out_ch, 1)); - ds = ggml_cont(ctx0, ggml_permute(ctx0, ds, 1, 2, 0, 3)); // conv→internal - ds = sam3_layer_norm_2d(ctx0, ds, model.mem_enc.ds_norm_w[s], model.mem_enc.ds_norm_b[s]); - ds = ggml_gelu(ctx0, ds); - ds = ggml_cont(ctx0, ggml_permute(ctx0, ds, 2, 0, 1, 3)); // internal→conv + // Match new PCS detections against ALL tracked objects (both active AND pending) + // so that the same object isn't detected as a new instance every frame. + std::vector all_tracked; + std::vector all_tracked_masks; + for (auto& ml : tracker.masklets) { + all_tracked.push_back(ml); + auto it = pm.find(ml.instance_id); + all_tracked_masks.push_back(it != pm.end() ? it->second : sam3_mask{}); } - // Final 1×1 conv: channels → D=256 - ds = ggml_conv_2d(ctx0, model.mem_enc.ds_conv_w[4], ds, 1, 1, 0, 0, 1, 1); - ds = ggml_add(ctx0, ds, ggml_reshape_4d(ctx0, model.mem_enc.ds_conv_b[4], 1, 1, D, 1)); - ds = ggml_cont(ctx0, ggml_permute(ctx0, ds, 1, 2, 0, 3)); // conv→internal [D, ...] + for (auto& ml : tracker.pending) { + all_tracked.push_back(ml); + auto it = pm.find(ml.instance_id); + all_tracked_masks.push_back(it != pm.end() ? it->second : sam3_mask{}); + } + auto mat = sam3_match_detections(all_tracked, nd.detections, + all_tracked_masks, tracker.params.assoc_iou_threshold); + std::vector dmat(nd.detections.size(), false); + for (auto& m : mat) { + dmat[m.second] = true; + // Update last_seen on the matched masklet (in active or pending) + int mid = all_tracked[m.first].instance_id; + for (auto& ml : tracker.masklets) + if (ml.instance_id == mid) { + ml.last_seen = fi; + break; + } + for (auto& ml : tracker.pending) + if (ml.instance_id == mid) { + ml.last_seen = fi; + ml.mds_sum++; + break; + } + } + // Map: pending instance_id → PCS detection index (for display) + std::map pending_det_idx; - // ── Pixel feature projection ── - auto* neck2_whcb = ggml_cont(ctx0, ggml_permute(ctx0, neck2, 2, 0, 1, 3)); - auto* pix = ggml_conv_2d(ctx0, model.mem_enc.pix_proj_w, neck2_whcb, 1, 1, 0, 0, 1, 1); - pix = ggml_add(ctx0, pix, ggml_reshape_4d(ctx0, model.mem_enc.pix_proj_b, 1, 1, D, 1)); - pix = ggml_cont(ctx0, ggml_permute(ctx0, pix, 1, 2, 0, 3)); - sam3_name_tensorf(pix, "phase7_mem%d_pix_proj", slot_idx); + for (int j = 0; j < (int)nd.detections.size(); ++j) { + if (dmat[j]) continue; + sam3_masklet ml; + ml.instance_id = tracker.next_inst_id++; + ml.first_frame = fi; + ml.last_seen = fi; + ml.last_score = nd.detections[j].score; + ml.mds_sum = 1; + pending_det_idx[ml.instance_id] = j; - // ── Fusion: ADD (not multiply) ── - auto* fused = ggml_add(ctx0, pix, ds); - sam3_name_tensorf(fused, "phase7_mem%d_fused_input", slot_idx); + // Encode the detection's mask into memory so subsequent frames can + // propagate from it. Without this, pending masklets have no memory + // and propagation returns empty on all future frames. + const auto& det = nd.detections[j]; + if (!det.mask.data.empty()) { + // Build mask logits from the PCS detection's binary mask. + const int mh = sam3_eff_feat_size(state, model.hparams) * 4, mw = mh; + // Resize to mh×mw logit space: inside mask → +5.0, outside → -5.0. + std::vector det_logits(mh * mw); + for (int r = 0; r < mh; ++r) { + for (int c = 0; c < mw; ++c) { + int sx = c * det.mask.width / mw; + int sy = r * det.mask.height / mh; + sx = std::min(sx, det.mask.width - 1); + sy = std::min(sy, det.mask.height - 1); + det_logits[r * mw + c] = (det.mask.data[sy * det.mask.width + sx] > 127) + ? 5.0f + : -5.0f; + } + } + sam3_encode_memory(tracker, state, model, ml.instance_id, + det_logits.data(), mh, mw, fi, true, 1.0f); - // ── CXBlocks ── - auto* h = fused; - struct ggml_tensor* fuser0 = nullptr; - struct ggml_tensor* fuser1 = nullptr; - for (int i = 0; i < 2; ++i) { - h = sam3_cxblock_forward(ctx0, h, - model.mem_enc.fuser_dw_w[i], model.mem_enc.fuser_dw_b[i], - model.mem_enc.fuser_norm_w[i], model.mem_enc.fuser_norm_b[i], - model.mem_enc.fuser_fc1_w[i], model.mem_enc.fuser_fc1_b[i], - model.mem_enc.fuser_fc2_w[i], model.mem_enc.fuser_fc2_b[i], - model.mem_enc.fuser_gamma[i]); - if (i == 0) { - sam3_name_tensorf(h, "phase7_mem%d_fuser0", slot_idx); - fuser0 = h; - } else { - sam3_name_tensorf(h, "phase7_mem%d_fuser1", slot_idx); - fuser1 = h; + // Store a dummy object pointer (from no_obj_ptr since we don't + // have a SAM token from PCS — PCS uses a different decoder) + std::vector no_ptr(D); + auto* nop = model.tensors.at("no_obj_ptr"); + if (nop->type == GGML_TYPE_F16) { + std::vector tmp(D); + ggml_backend_tensor_get(nop, tmp.data(), 0, D * sizeof(ggml_fp16_t)); + ggml_fp16_to_fp32_row(tmp.data(), no_ptr.data(), D); + } else { + ggml_backend_tensor_get(nop, no_ptr.data(), 0, D * sizeof(float)); + } + sam3_store_obj_ptr(tracker, model, ml.instance_id, no_ptr.data(), fi); } + + tracker.pending.push_back(std::move(ml)); + } + for (auto& ml : tracker.masklets) { + int id = ml.instance_id; + auto it = po.find(id); + if (it == po.end() || it->second.mask_logits.empty()) continue; + sam3_encode_memory(tracker, state, model, id, + it->second.mask_logits.data(), it->second.mask_h, it->second.mask_w, fi, false, it->second.obj_score); + std::vector op(D); + sam3_extract_obj_ptr_cpu(model, it->second.sam_token.data(), it->second.obj_score, op.data()); + sam3_store_obj_ptr(tracker, model, id, op.data(), fi); } + sam3_update_tracker(tracker, fi); - // ── Output projection ── - auto* h_whcb = ggml_cont(ctx0, ggml_permute(ctx0, h, 2, 0, 1, 3)); - auto* mo = ggml_conv_2d(ctx0, model.mem_enc.out_proj_w, h_whcb, 1, 1, 0, 0, 1, 1); - mo = ggml_add(ctx0, mo, ggml_reshape_4d(ctx0, model.mem_enc.out_proj_b, 1, 1, MD, 1)); - mo = ggml_cont(ctx0, ggml_permute(ctx0, mo, 1, 2, 0, 3)); - sam3_name_tensorf(mo, "phase7_mem%d_output", slot_idx); - - ggml_set_output(pix); - ggml_set_output(fused); - if (fuser0) ggml_set_output(fuser0); - if (fuser1) ggml_set_output(fuser1); - ggml_set_output(mo); - - auto* graph = ggml_new_graph_custom(ctx0, 16384, false); - ggml_build_forward_expand(graph, mo); - ggml_build_forward_expand(graph, pix); - ggml_build_forward_expand(graph, fused); - if (fuser0) ggml_build_forward_expand(graph, fuser0); - if (fuser1) ggml_build_forward_expand(graph, fuser1); + // Helper: build detection from a mask and add to result + auto add_mask_to_result = [&](int inst_id, float score, const sam3_mask& mask) { + if (mask.data.empty()) return; + sam3_detection det; + det.instance_id = inst_id; + det.score = score; + det.mask = mask; + det.mask.instance_id = inst_id; + det.mask.iou_score = score; + float x0 = 1e9f, y0 = 1e9f, x1 = -1e9f, y1 = -1e9f; + for (int p = 0; p < (int)det.mask.data.size(); ++p) + if (det.mask.data[p] > 127) { + int x = p % det.mask.width, y = p / det.mask.width; + x0 = std::min(x0, (float)x); + y0 = std::min(y0, (float)y); + x1 = std::max(x1, (float)x); + y1 = std::max(y1, (float)y); + } + if (x0 <= x1) det.box = {x0, y0, x1, y1}; + result.detections.push_back(std::move(det)); + }; - auto* galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); - if (!ggml_gallocr_reserve(galloc, graph) || !ggml_gallocr_alloc_graph(galloc, graph)) { - fprintf(stderr, "%s: failed to allocate phase 7 memory graph\n", __func__); - ggml_gallocr_free(galloc); - ggml_free(ctx0); - return false; + // Include confirmed (active) masklet propagation results + for (auto& ml : tracker.masklets) { + auto it = pm.find(ml.instance_id); + if (it != pm.end()) add_mask_to_result(ml.instance_id, ml.last_score, it->second); } - ggml_backend_tensor_set(neck2, neck2_data.data(), 0, neck2_data.size() * sizeof(float)); - // Mask is in WHCB layout = row-major over (w, h): just upload flat - ggml_backend_tensor_set(mask_in, m_interp.data(), 0, m_interp.size() * sizeof(float)); - - sam3_graph_compute(model.backend, graph, n_threads); - - // Read output and apply no_obj_embed_spatial if occluded - mem_out_data.resize(MD * H * H); - ggml_backend_tensor_get(mo, mem_out_data.data(), 0, mem_out_data.size() * sizeof(float)); - - if (obj_score <= 0.0f && model.no_obj_embed_spatial) { - std::vector no_obj_emb(MD); - auto* noe = model.no_obj_embed_spatial; - if (noe->type == GGML_TYPE_F16) { - std::vector tmp(MD); - ggml_backend_tensor_get(noe, tmp.data(), 0, MD * sizeof(ggml_fp16_t)); - ggml_fp16_to_fp32_row(tmp.data(), no_obj_emb.data(), MD); + // Include pending masklet results: on the detection frame, use the PCS + // detection mask; on subsequent frames, use the propagated mask. + for (auto& ml : tracker.pending) { + auto it = pm.find(ml.instance_id); + if (it != pm.end() && !it->second.data.empty()) { + // Propagated mask available (frame > detection frame) + add_mask_to_result(ml.instance_id, ml.last_score, it->second); } else { - ggml_backend_tensor_get(noe, no_obj_emb.data(), 0, MD * sizeof(float)); + // First detection frame: use the PCS detection's mask+box directly + auto di = pending_det_idx.find(ml.instance_id); + if (di != pending_det_idx.end() && di->second < (int)nd.detections.size()) { + const auto& det = nd.detections[di->second]; + if (!det.mask.data.empty()) { + add_mask_to_result(ml.instance_id, ml.last_score, det.mask); + } + } } - for (int i = 0; i < MD * H * H; ++i) - mem_out_data[i] += no_obj_emb[i % MD]; } + sam3_resolve_overlaps(result.detections); + for (auto& d : result.detections) { + if (d.mask.data.empty()) continue; + sam3_fill_holes(d.mask.data.data(), d.mask.width, d.mask.height, tracker.params.fill_hole_area); + sam3_remove_sprinkles(d.mask.data.data(), d.mask.width, d.mask.height, tracker.params.fill_hole_area); + } + tracker.frame_index++; + SAM3_LOG(2, "%s: frame %d done — %zu tracked\n", __func__, fi, result.detections.size()); + return result; +} - // Compute sinusoidal spatial PE [MD, 72, 72] - mem_pe_data = sam3_sinusoidal_pe_2d(H, H, MD); - - bool ok = true; - const std::vector tensor_names = { - "phase7_mem" + std::to_string(slot_idx) + "_pix_proj", - "phase7_mem" + std::to_string(slot_idx) + "_fused_input", - "phase7_mem" + std::to_string(slot_idx) + "_fuser0", - "phase7_mem" + std::to_string(slot_idx) + "_fuser1", - "phase7_mem" + std::to_string(slot_idx) + "_output", - }; - for (const auto& name : tensor_names) { - auto* t = ggml_get_tensor(ctx0, name.c_str()); - if (!t || !sam3_dump_tensor_to_path(t, name, output_dir + "/" + name)) { - ok = false; +bool sam3_refine_instance(sam3_tracker& tracker, sam3_state& state, + const sam3_model& model, int instance_id, + const std::vector& pos_points, + const std::vector& neg_points) { + const int D = model.hparams.neck_dim; + sam3_masklet* tgt = nullptr; + for (auto& ml : tracker.masklets) + if (ml.instance_id == instance_id) { + tgt = &ml; + break; } + if (!tgt) + for (auto& ml : tracker.pending) + if (ml.instance_id == instance_id) { + tgt = &ml; + break; + } + if (!tgt) { + fprintf(stderr, "%s: instance %d not found\n", __func__, instance_id); + return false; } - // Also dump the output with no_obj_embed_spatial applied - if (!sam3_dump_raw_f32_to_path(mem_out_data.data(), {MD, H, H, 1}, - output_dir + "/phase7_mem" + std::to_string(slot_idx) + "_output")) { - ok = false; + sam3_pvs_params pvs; + pvs.pos_points = pos_points; + pvs.neg_points = neg_points; + pvs.multimask = false; + auto r = sam3_segment_pvs(state, model, pvs); + if (r.detections.empty()) return false; + // tracker.frame_index points to the *next* frame; the refinement applies + // to the frame that was last tracked / encoded. + int fi = std::max(0, tracker.frame_index - 1); + const auto& rdet = r.detections[0]; + tgt->last_score = rdet.score; + tgt->last_seen = fi; + std::vector op(D); + if (!rdet.sam_token.empty()) { + sam3_extract_obj_ptr_cpu(model, rdet.sam_token.data(), rdet.mask.obj_score, op.data()); + } else { + std::fill(op.begin(), op.end(), 0.0f); } - - ggml_gallocr_free(galloc); - ggml_free(ctx0); - return ok; + sam3_store_obj_ptr(tracker, model, instance_id, op.data(), fi); + SAM3_LOG(2, "%s: refined instance %d\n", __func__, instance_id); + return true; } -bool sam3_test_dump_phase7_from_ref_inputs(const sam3_model& model, - const std::string& case_ref_dir, - const std::string& output_dir, - int n_threads) { - const auto& hp = model.hparams; - const int D = hp.neck_dim; - const int MD = hp.mem_out_dim; - const int H = hp.feat_size(); - const int N = H * H; - const int H1 = H * 2; - const int H0 = H * 4; - const int MASK_HW = H * 4; - const int half_d = D / 2; // 128 +int sam3_tracker_add_instance(sam3_tracker& tracker, sam3_state& state, + const sam3_model& model, + const sam3_pvs_params& pvs_params) { + const int D = model.hparams.neck_dim; + const int mask_hw = sam3_eff_feat_size(state, model.hparams) * 4; - std::map meta; - if (!sam3_load_kv_text_file(case_ref_dir + "/meta.txt", meta)) { - return false; - } - const int num_slots = sam3_meta_get_int(meta, "num_slots", 0); - if (num_slots <= 0) { - fprintf(stderr, "%s: no memory slots in %s\n", __func__, case_ref_dir.c_str()); - return false; + // Run PVS to get the segmentation mask + auto r = sam3_segment_pvs(state, model, pvs_params); + if (r.detections.empty()) { + fprintf(stderr, "%s: PVS returned no masks\n", __func__); + return -1; } - std::vector neck_trk_0_nchw; - std::vector neck_trk_1_nchw; - std::vector neck_trk_2_nchw; - if (!sam3_load_ref_f32_data(case_ref_dir + "/input_curr_neck_trk_0", neck_trk_0_nchw, D * H0 * H0) || - !sam3_load_ref_f32_data(case_ref_dir + "/input_curr_neck_trk_1", neck_trk_1_nchw, D * H1 * H1) || - !sam3_load_ref_f32_data(case_ref_dir + "/input_curr_neck_trk_2", neck_trk_2_nchw, D * H * H)) { - return false; + const auto& det = r.detections[0]; + if (det.mask.data.empty()) { + fprintf(stderr, "%s: PVS mask is empty\n", __func__); + return -1; } - auto neck_trk_0 = sam3_reorder_nchw_to_ggml_dwh(neck_trk_0_nchw, D, H0, H0); - auto neck_trk_1 = sam3_reorder_nchw_to_ggml_dwh(neck_trk_1_nchw, D, H1, H1); - auto neck_trk_2 = sam3_reorder_nchw_to_ggml_dwh(neck_trk_2_nchw, D, H, H); - - // ── Memory encoding per slot ──────────────────────────────────────── - std::vector> mem_slot_outputs(num_slots); - std::vector> mem_slot_pes(num_slots); - std::vector> obj_ptrs; - std::vector slot_spatial_tpos(num_slots); - std::vector slot_ptr_tpos(num_slots); - bool ok = true; - - for (int i = 0; i < num_slots; ++i) { - std::vector mask_logits_nchw; - std::vector sam_token; - std::vector obj_score; - if (!sam3_load_ref_f32_data(case_ref_dir + "/input_mem_mask_logits_" + std::to_string(i), - mask_logits_nchw, MASK_HW * MASK_HW) || - !sam3_load_ref_f32_data(case_ref_dir + "/input_mem_sam_token_" + std::to_string(i), - sam_token, D) || - !sam3_load_ref_f32_data(case_ref_dir + "/input_mem_obj_score_" + std::to_string(i), - obj_score, 1)) { - return false; - } - - slot_spatial_tpos[i] = sam3_meta_get_int(meta, "slot" + std::to_string(i) + "_spatial_tpos", 0); - slot_ptr_tpos[i] = sam3_meta_get_int(meta, "slot" + std::to_string(i) + "_ptr_tpos", -1); - - if (!sam3_test_dump_phase7_mem_slot_current(model, neck_trk_2, mask_logits_nchw, - obj_score[0], i, output_dir, - mem_slot_outputs[i], mem_slot_pes[i], n_threads)) { - ok = false; - } + int inst_id = tracker.next_inst_id++; + // tracker.frame_index points to the *next* frame to process; + // the instance is being added on the frame that was just tracked. + int fi = tracker.frame_index - 1; + if (fi < 0) fi = 0; - std::vector obj_ptr(D, 0.0f); - sam3_extract_obj_ptr_cpu(model, sam_token.data(), obj_score[0], obj_ptr.data()); - if (!sam3_dump_raw_f32_to_path(obj_ptr.data(), {D, 1}, - output_dir + "/phase7_obj_ptr" + std::to_string(i))) { - ok = false; + // Create synthetic 288x288 logits from the binary mask. + // sam3_encode_memory applies sigmoid then scale/bias, so +6/-6 gives + // sigmoid values ~0.9975/~0.0025 — practically identical to real logits. + std::vector synth_logits(mask_hw * mask_hw); + { + int mw = det.mask.width, mh = det.mask.height; + for (int y = 0; y < mask_hw; ++y) { + int sy = y * mh / mask_hw; + for (int x = 0; x < mask_hw; ++x) { + int sx = x * mw / mask_hw; + synth_logits[y * mask_hw + x] = + (det.mask.data[sy * mw + sx] > 127) ? 6.0f : -6.0f; + } } - - obj_ptrs.push_back(obj_ptr); - } - - // ── Build prompt and prompt_pos ───────────────────────────────────── - std::vector ptr_tpos_vec(num_slots); - for (int i = 0; i < num_slots; ++i) - ptr_tpos_vec[i] = slot_ptr_tpos[i]; - - auto pd = sam3_build_prompt_and_pos(model, mem_slot_outputs, mem_slot_pes, - slot_spatial_tpos, obj_ptrs, ptr_tpos_vec); - - // ── Precompute RoPE frequencies ───────────────────────────────────── - std::vector rope_q_data(N * D); - sam3_compute_axial_cis(rope_q_data.data(), D, H, H, 10000.0f, 1.0f); - - std::vector rope_k_data; - if (pd.M_spatial > 0) { - int n_spatial_slots = pd.M_spatial / N; - rope_k_data.resize(2 * half_d * pd.M_spatial); - for (int s = 0; s < n_spatial_slots; ++s) - memcpy(rope_k_data.data() + s * D * N, rope_q_data.data(), D * N * sizeof(float)); } - // ── Build graph ───────────────────────────────────────────────────── - const size_t buf_size = ggml_tensor_overhead() * 32768 + ggml_graph_overhead() * 2; - struct ggml_init_params gparams = {buf_size, nullptr, true}; - struct ggml_context* ctx0 = ggml_init(gparams); - if (!ctx0) { - fprintf(stderr, "%s: failed to init phase 7 context\n", __func__); - return false; + // Encode into memory bank + float obj_score = det.mask.obj_score; + if (!sam3_encode_memory(tracker, state, model, inst_id, + synth_logits.data(), mask_hw, mask_hw, + fi, true, obj_score)) { + fprintf(stderr, "%s: failed to encode memory for instance %d\n", __func__, inst_id); + return -1; } - auto* neck0 = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, D, H0, H0, 1); - auto* neck1 = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, D, H1, H1, 1); - auto* neck2 = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, D, H, H, 1); - ggml_set_name(neck0, "phase7_curr_neck0"); - ggml_set_name(neck1, "phase7_curr_neck1"); - ggml_set_name(neck2, "phase7_curr_neck2"); - ggml_set_input(neck0); - ggml_set_input(neck1); - ggml_set_input(neck2); - - auto* curr = ggml_reshape_3d(ctx0, neck2, D, N, 1); - - // src_pos for pos_enc_at_input - auto* src_pos_t = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, D, N, 1); - ggml_set_name(src_pos_t, "phase7_src_pos"); - ggml_set_input(src_pos_t); - - // Prompt and prompt_pos - auto* prompt_t = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, MD, pd.M_total, 1); - ggml_set_name(prompt_t, "phase7_prompt"); - ggml_set_input(prompt_t); - auto* prompt_pos_t = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, MD, pd.M_total, 1); - ggml_set_name(prompt_pos_t, "phase7_prompt_pos"); - ggml_set_input(prompt_pos_t); - - // RoPE frequencies - auto* rope_q_t = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 2, half_d, N); - ggml_set_name(rope_q_t, "phase7_rope_q"); - ggml_set_input(rope_q_t); - struct ggml_tensor* rope_k_t = nullptr; - if (pd.M_spatial > 0) { - rope_k_t = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 2, half_d, pd.M_spatial); - ggml_set_name(rope_k_t, "phase7_rope_k"); - ggml_set_input(rope_k_t); + // Extract object pointer from the SAM decoder token + std::vector op(D); + if (!det.sam_token.empty()) { + sam3_extract_obj_ptr_cpu(model, det.sam_token.data(), obj_score, op.data()); + } else { + std::fill(op.begin(), op.end(), 0.0f); } + sam3_store_obj_ptr(tracker, model, inst_id, op.data(), fi); - auto* conditioned = sam3_build_mem_attn_graph(ctx0, model, curr, src_pos_t, - prompt_t, prompt_pos_t, - rope_q_t, rope_k_t, - pd.num_obj_ptr_tokens); - auto* cond_spatial = ggml_reshape_4d(ctx0, conditioned, D, H, H, 1); - - // Bug 3 fix: single not_a_point_embed token - auto* sparse_in = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, D, 1, 1); - ggml_set_name(sparse_in, "phase7_sparse"); - ggml_set_input(sparse_in); - - auto* image_pe = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, D, H, H, 1); - auto* dense_emb = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, D, H, H, 1); - ggml_set_name(image_pe, "phase7_image_pe"); - ggml_set_name(dense_emb, "phase7_dense"); - ggml_set_input(image_pe); - ggml_set_input(dense_emb); + // Create confirmed masklet + sam3_masklet ml; + ml.instance_id = inst_id; + ml.first_frame = fi; + ml.last_seen = fi; + ml.last_score = det.score; + ml.confirmed = true; + ml.mds_sum = 1; + tracker.masklets.push_back(std::move(ml)); - auto dec = sam3_build_sam_dec_graph(ctx0, model, cond_spatial, image_pe, - sparse_in, dense_emb, neck0, neck1); - - auto* mask0 = ggml_view_3d(ctx0, dec.masks, MASK_HW * MASK_HW, 1, 1, - dec.masks->nb[1], dec.masks->nb[2], 0); - mask0 = ggml_cont(ctx0, mask0); - ggml_set_name(mask0, "phase7_prop_masks"); - - auto* iou0 = ggml_view_2d(ctx0, dec.iou_pred, 1, 1, dec.iou_pred->nb[1], 0); - iou0 = ggml_cont(ctx0, iou0); - ggml_set_name(iou0, "phase7_prop_iou"); - ggml_set_name(dec.obj_score, "phase7_prop_obj_score"); - ggml_set_name(dec.sam_token, "phase7_prop_sam_token"); - - ggml_set_output(conditioned); - ggml_set_output(mask0); - ggml_set_output(iou0); - ggml_set_output(dec.obj_score); - ggml_set_output(dec.sam_token); + SAM3_LOG(2, "%s: added instance #%d (score=%.3f)\n", __func__, inst_id, det.score); + return inst_id; +} - const std::vector mem_attn_names = { - "phase7_mem_attn_input", - "phase7_mem_attn_layer0_after_sa", - "phase7_mem_attn_layer0_after_ca", - "phase7_mem_attn_layer0_after_ffn", - "phase7_mem_attn_layer1_after_sa", - "phase7_mem_attn_layer1_after_ca", - "phase7_mem_attn_layer1_after_ffn", - "phase7_mem_attn_layer2_after_sa", - "phase7_mem_attn_layer2_after_ca", - "phase7_mem_attn_layer2_after_ffn", - "phase7_mem_attn_layer3_after_sa", - "phase7_mem_attn_layer3_after_ca", - "phase7_mem_attn_layer3_after_ffn", - "phase7_mem_attn_output", - "phase7_prop_masks", - "phase7_prop_iou", - "phase7_prop_obj_score", - "phase7_prop_sam_token", - }; - for (const auto& name : mem_attn_names) { - auto* t = ggml_get_tensor(ctx0, name.c_str()); - if (t) ggml_set_output(t); - } +int sam3_tracker_frame_index(const sam3_tracker& tracker) { return tracker.frame_index; } - auto* graph = ggml_new_graph_custom(ctx0, 32768, false); - ggml_build_forward_expand(graph, mask0); - ggml_build_forward_expand(graph, iou0); - ggml_build_forward_expand(graph, dec.obj_score); - ggml_build_forward_expand(graph, dec.sam_token); - ggml_build_forward_expand(graph, conditioned); - for (const auto& name : mem_attn_names) { - auto* t = ggml_get_tensor(ctx0, name.c_str()); - if (t) ggml_build_forward_expand(graph, t); +void sam3_tracker_reset(sam3_tracker& tracker) { + tracker.frame_index = 0; + tracker.next_inst_id = 1; + tracker.masklets.clear(); + tracker.pending.clear(); + tracker.mem_banks.clear(); + tracker.ptr_banks.clear(); + for (auto* b : tracker.owned_buffers) + if (b) ggml_backend_buffer_free(b); + tracker.owned_buffers.clear(); + if (tracker.ctx) { + ggml_free(tracker.ctx); + tracker.ctx = nullptr; } - - auto* galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend)); - if (!ggml_gallocr_reserve(galloc, graph) || !ggml_gallocr_alloc_graph(galloc, graph)) { - fprintf(stderr, "%s: failed to allocate phase 7 graph\n", __func__); - ggml_gallocr_free(galloc); - ggml_free(ctx0); - return false; + if (tracker.buffer) { + ggml_backend_buffer_free(tracker.buffer); + tracker.buffer = nullptr; } +} - // Upload neck features - ggml_backend_tensor_set(neck0, neck_trk_0.data(), 0, neck_trk_0.size() * sizeof(float)); - ggml_backend_tensor_set(neck1, neck_trk_1.data(), 0, neck_trk_1.size() * sizeof(float)); - ggml_backend_tensor_set(neck2, neck_trk_2.data(), 0, neck_trk_2.size() * sizeof(float)); +/***************************************************************************** +** Visual-only video tracking +*****************************************************************************/ - // Upload prompt and prompt_pos - ggml_backend_tensor_set(prompt_t, pd.prompt.data(), 0, pd.prompt.size() * sizeof(float)); - ggml_backend_tensor_set(prompt_pos_t, pd.prompt_pos.data(), 0, pd.prompt_pos.size() * sizeof(float)); +sam3_tracker_ptr sam3_create_visual_tracker( + const sam3_model& model, + const sam3_visual_track_params& params) { + sam3_video_params vp; + vp.text_prompt = ""; // no PCS detection + vp.assoc_iou_threshold = params.assoc_iou_threshold; + vp.max_keep_alive = params.max_keep_alive; + vp.recondition_every = params.recondition_every; + vp.fill_hole_area = params.fill_hole_area; + sam3_tracker_ptr tracker(new sam3_tracker()); + tracker->params = vp; + fprintf(stderr, "%s: visual-only tracker created (max_keep_alive=%d)\n", + __func__, params.max_keep_alive); + return tracker; +} - // Upload RoPE frequencies - ggml_backend_tensor_set(rope_q_t, rope_q_data.data(), 0, rope_q_data.size() * sizeof(float)); - if (rope_k_t && !rope_k_data.empty()) - ggml_backend_tensor_set(rope_k_t, rope_k_data.data(), 0, rope_k_data.size() * sizeof(float)); +sam3_result sam3_propagate_frame( + sam3_tracker& tracker, sam3_state& state, + const sam3_model& model, const sam3_image& frame) { + sam3_result result; + const int D = model.hparams.neck_dim; + if (!sam3_encode_image(state, model, frame)) return result; + int fi = tracker.frame_index; + fprintf(stderr, "%s: frame %d (%zu active + %zu pending)\n", + __func__, fi, tracker.masklets.size(), tracker.pending.size()); - // Upload src_pos (sinusoidal PE 256-dim for 72×72) - auto src_pos_data = sam3_sinusoidal_pe_2d(H, H, D); - ggml_backend_tensor_set(src_pos_t, src_pos_data.data(), 0, src_pos_data.size() * sizeof(float)); + // ── Propagate active masklets ──────────────────────────────────────── + std::map pm; + std::map po; + for (auto& ml : tracker.masklets) { + int id = ml.instance_id; + auto im = tracker.mem_banks.find(id); + if (im == tracker.mem_banks.end() || im->second.empty()) continue; + po[id] = sam3_propagate_single(tracker, state, model, ml, im->second, tracker.ptr_banks[id]); + if (po[id].mask_logits.empty()) continue; + auto rs = sam3_bilinear_interpolate(po[id].mask_logits.data(), + po[id].mask_w, po[id].mask_h, + state.orig_width, state.orig_height); + pm[id].width = state.orig_width; + pm[id].height = state.orig_height; + pm[id].data.resize(state.orig_width * state.orig_height); + int fg = 0; + for (int p = 0; p < (int)rs.size(); ++p) { + bool f = rs[p] > 0.0f; + pm[id].data[p] = f ? 255 : 0; + if (f) fg++; + } + ml.last_score = po[id].iou_scores[0]; + ml.last_seen = fi; + float cov = (float)fg / (state.orig_width * state.orig_height); + ml.mds_sum += (cov > 0.001f && po[id].obj_score > 0.0f) ? 1 : -1; + } - // Upload not_a_point_embed for sparse prompt - { - float nap[256]; - if (model.sam_pe.not_a_point_embed->type == GGML_TYPE_F16) { - std::vector tmp(D); - ggml_backend_tensor_get(model.sam_pe.not_a_point_embed, tmp.data(), 0, D * sizeof(ggml_fp16_t)); - ggml_fp16_to_fp32_row(tmp.data(), nap, D); - } else { - ggml_backend_tensor_get(model.sam_pe.not_a_point_embed, nap, 0, D * sizeof(float)); + // ── Propagate pending masklets ─────────────────────────────────────── + for (auto& ml : tracker.pending) { + int id = ml.instance_id; + auto im = tracker.mem_banks.find(id); + if (im == tracker.mem_banks.end() || im->second.empty()) continue; + auto p2 = sam3_propagate_single(tracker, state, model, ml, im->second, tracker.ptr_banks[id]); + if (!p2.mask_logits.empty()) { + ml.last_score = p2.iou_scores[0]; + ml.last_seen = fi; + auto r2 = sam3_bilinear_interpolate(p2.mask_logits.data(), + p2.mask_w, p2.mask_h, + state.orig_width, state.orig_height); + int fg2 = 0; + for (auto v : r2) + if (v > 0.0f) fg2++; + float c2 = (float)fg2 / (state.orig_width * state.orig_height); + ml.mds_sum += (c2 > 0.001f && p2.obj_score > 0.0f) ? 1 : -1; + pm[id].width = state.orig_width; + pm[id].height = state.orig_height; + pm[id].data.resize(state.orig_width * state.orig_height); + for (int p = 0; p < (int)r2.size(); ++p) + pm[id].data[p] = r2[p] > 0.0f ? 255 : 0; + sam3_encode_memory(tracker, state, model, id, + p2.mask_logits.data(), p2.mask_h, p2.mask_w, + fi, false, p2.obj_score); + std::vector op(D); + sam3_extract_obj_ptr_cpu(model, p2.sam_token.data(), p2.obj_score, op.data()); + sam3_store_obj_ptr(tracker, model, id, op.data(), fi); } - ggml_backend_tensor_set(sparse_in, nap, 0, D * sizeof(float)); } - // Upload image_pe and dense_emb - sam3_state pe_cache_state = {}; - sam3_populate_pe_cache(pe_cache_state, model); - ggml_backend_tensor_set(image_pe, pe_cache_state.dense_pe_cache.data(), 0, - pe_cache_state.dense_pe_cache.size() * sizeof(float)); - ggml_backend_tensor_set(dense_emb, pe_cache_state.dense_nomask_cache.data(), 0, - pe_cache_state.dense_nomask_cache.size() * sizeof(float)); + // ── Encode memory for active masklets ──────────────────────────────── + for (auto& ml : tracker.masklets) { + int id = ml.instance_id; + auto it = po.find(id); + if (it == po.end() || it->second.mask_logits.empty()) continue; + sam3_encode_memory(tracker, state, model, id, + it->second.mask_logits.data(), it->second.mask_h, + it->second.mask_w, fi, false, it->second.obj_score); + std::vector op(D); + sam3_extract_obj_ptr_cpu(model, it->second.sam_token.data(), + it->second.obj_score, op.data()); + sam3_store_obj_ptr(tracker, model, id, op.data(), fi); + } - sam3_graph_compute(model.backend, graph, n_threads); + // ── Update tracker state (confirmation / eviction) ─────────────────── + sam3_update_tracker(tracker, fi); - // Dump the mem_attn_input manually (it's now x = curr + 0.1*src_pos, not just curr) - // The named tensor "phase7_mem_attn_input" is no longer in the graph since we - // removed the explicit naming. But all the per-layer outputs are still named. - for (const auto& name : mem_attn_names) { - auto* t = ggml_get_tensor(ctx0, name.c_str()); - if (!t || !sam3_dump_tensor_to_path(t, name, output_dir + "/" + name)) { - ok = false; - } + // ── Build result ───────────────────────────────────────────────────── + auto add_mask_to_result = [&](int inst_id, float score, const sam3_mask& mask) { + if (mask.data.empty()) return; + sam3_detection det; + det.instance_id = inst_id; + det.score = score; + det.mask = mask; + det.mask.instance_id = inst_id; + det.mask.iou_score = score; + float x0 = 1e9f, y0 = 1e9f, x1 = -1e9f, y1 = -1e9f; + for (int p = 0; p < (int)det.mask.data.size(); ++p) + if (det.mask.data[p] > 127) { + int x = p % det.mask.width, y = p / det.mask.width; + x0 = std::min(x0, (float)x); + y0 = std::min(y0, (float)y); + x1 = std::max(x1, (float)x); + y1 = std::max(y1, (float)y); + } + if (x0 <= x1) det.box = {x0, y0, x1, y1}; + result.detections.push_back(std::move(det)); + }; + + for (auto& ml : tracker.masklets) { + auto it = pm.find(ml.instance_id); + if (it != pm.end()) add_mask_to_result(ml.instance_id, ml.last_score, it->second); + } + for (auto& ml : tracker.pending) { + auto it = pm.find(ml.instance_id); + if (it != pm.end()) add_mask_to_result(ml.instance_id, ml.last_score, it->second); } - ggml_gallocr_free(galloc); - ggml_free(ctx0); - return ok; + sam3_resolve_overlaps(result.detections); + for (auto& d : result.detections) { + if (d.mask.data.empty()) continue; + sam3_fill_holes(d.mask.data.data(), d.mask.width, d.mask.height, + tracker.params.fill_hole_area); + sam3_remove_sprinkles(d.mask.data.data(), d.mask.width, d.mask.height, + tracker.params.fill_hole_area); + } + tracker.frame_index++; + SAM3_LOG(2, "%s: frame %d done — %zu tracked\n", + __func__, fi, result.detections.size()); + return result; } /***************************************************************************** -** Debug: dump state tensors +** Utility — image I/O *****************************************************************************/ -static struct ggml_tensor * sam3_find_state_tensor(const sam3_state & state, - const std::string & tensor_name) { - struct ggml_tensor* t = nullptr; - - if (tensor_name == "vit_output") { - t = state.vit_output; - } else if (tensor_name == "neck_det_0") { - t = state.neck_det[0]; - } else if (tensor_name == "neck_det_1") { - t = state.neck_det[1]; - } else if (tensor_name == "neck_det_2") { - t = state.neck_det[2]; - } else if (tensor_name == "neck_det_3") { - t = state.neck_det[3]; - } else if (tensor_name == "neck_trk_0") { - t = state.neck_trk[0]; - } else if (tensor_name == "neck_trk_1") { - t = state.neck_trk[1]; - } else if (tensor_name == "neck_trk_2") { - t = state.neck_trk[2]; - } else if (tensor_name == "neck_trk_3") { - t = state.neck_trk[3]; - } else if (tensor_name == "neck_det_pe_0") { - t = state.neck_det_pe[0]; - } else if (tensor_name == "neck_det_pe_1") { - t = state.neck_det_pe[1]; - } else if (tensor_name == "neck_det_pe_2") { - t = state.neck_det_pe[2]; - } else if (tensor_name == "neck_det_pe_3") { - t = state.neck_det_pe[3]; - } else { - // Search by ggml name in the context - if (state.ctx) { - t = ggml_get_tensor(state.ctx, tensor_name.c_str()); - } - // Also search PE context - if (!t && state.pe_ctx) { - t = ggml_get_tensor(state.pe_ctx, tensor_name.c_str()); - } +sam3_image sam3_load_image(const std::string& path) { + sam3_image img; + int w, h, c; + uint8_t* data = stbi_load(path.c_str(), &w, &h, &c, 3); + if (!data) { + fprintf(stderr, "%s: failed to load '%s'\n", __func__, path.c_str()); + return img; } + img.width = w; + img.height = h; + img.channels = 3; + img.data.assign(data, data + w * h * 3); + stbi_image_free(data); + return img; +} - return t; +bool sam3_save_mask(const sam3_mask& mask, const std::string& path) { + if (mask.data.empty()) return false; + return stbi_write_png(path.c_str(), mask.width, mask.height, 1, + mask.data.data(), mask.width) != 0; } -static bool sam3_fill_tensor_info(struct ggml_tensor * t, sam3_tensor_info & info) { - if (!t) { - return false; - } +sam3_image sam3_decode_video_frame(const std::string& video_path, int frame_index) { + sam3_image img; - for (int i = 0; i < 4; ++i) { - info.ne[i] = t->ne[i]; - info.nb[i] = t->nb[i]; - } - info.type = (int) t->type; - info.op = (int) t->op; - info.is_contiguous = ggml_is_contiguous(t); - return true; -} + // Use ffmpeg to extract a single frame as raw RGB (frame-accurate) + char cmd[1024]; + snprintf(cmd, sizeof(cmd), + "ffmpeg -nostdin -loglevel error -i \"%s\" " + "-vf \"select=eq(n\\,%d)\" -vsync vfr -frames:v 1 " + "-f rawvideo -pix_fmt rgb24 pipe:1 2>%s", + video_path.c_str(), frame_index, SAM3_NULL_DEV); -bool sam3_get_state_tensor_info(const sam3_state & state, - const std::string & tensor_name, - sam3_tensor_info & info) { - struct ggml_tensor * t = sam3_find_state_tensor(state, tensor_name); - if (!t) { - fprintf(stderr, "%s: tensor '%s' not found in state\n", __func__, tensor_name.c_str()); - return false; + // First, get dimensions + char info_cmd[1024]; + snprintf(info_cmd, sizeof(info_cmd), + "ffprobe -v error -select_streams v:0 " + "-show_entries stream=width,height -of csv=p=0 \"%s\" 2>%s", + video_path.c_str(), SAM3_NULL_DEV); + FILE* fp = popen(info_cmd, SAM3_POPEN_READ); + if (!fp) return img; + int w = 0, h = 0; + if (fscanf(fp, "%d,%d", &w, &h) != 2) { + pclose(fp); + return img; } - return sam3_fill_tensor_info(t, info); -} + pclose(fp); -bool sam3_dump_state_tensor(const sam3_state& state, - const std::string& tensor_name, - const std::string& output_path) { - struct ggml_tensor * t = sam3_find_state_tensor(state, tensor_name); - if (!t) { - fprintf(stderr, "%s: tensor '%s' not found in state\n", __func__, tensor_name.c_str()); - return false; - } - return sam3_dump_tensor_to_path(t, tensor_name, output_path); -} + img.width = w; + img.height = h; + img.channels = 3; + img.data.resize(w * h * 3); -bool sam3_get_model_tensor_info(const sam3_model & model, - const std::string & tensor_name, - sam3_tensor_info & info) { - auto it = model.tensors.find(tensor_name); - if (it == model.tensors.end() || !it->second) { - fprintf(stderr, "%s: tensor '%s' not found in model\n", __func__, tensor_name.c_str()); - return false; + fp = popen(cmd, SAM3_POPEN_READ); + if (!fp) { + img.data.clear(); + return img; + } + size_t nread = fread(img.data.data(), 1, img.data.size(), fp); + pclose(fp); + if (nread != img.data.size()) { + img.data.clear(); } - return sam3_fill_tensor_info(it->second, info); + + return img; } -bool sam3_dump_model_tensor(const sam3_model & model, - const std::string & tensor_name, - const std::string & output_path) { - auto it = model.tensors.find(tensor_name); - if (it == model.tensors.end() || !it->second) { - fprintf(stderr, "%s: tensor '%s' not found in model\n", __func__, tensor_name.c_str()); - return false; +sam3_video_info sam3_get_video_info(const std::string& video_path) { + sam3_video_info info; + + char cmd[1024]; + snprintf(cmd, sizeof(cmd), + "ffprobe -v error -select_streams v:0 " + "-show_entries stream=width,height,r_frame_rate,nb_frames " + "-of csv=p=0 \"%s\" 2>%s", + video_path.c_str(), SAM3_NULL_DEV); + FILE* fp = popen(cmd, SAM3_POPEN_READ); + if (!fp) return info; + + int w = 0, h = 0, num = 0, den = 1, nf = 0; + if (fscanf(fp, "%d,%d,%d/%d,%d", &w, &h, &num, &den, &nf) >= 4) { + info.width = w; + info.height = h; + info.fps = (den > 0) ? static_cast(num) / den : 0.0f; + info.n_frames = nf; } - return sam3_dump_tensor_to_path(it->second, tensor_name, output_path); + pclose(fp); + return info; } diff --git a/sam3.h b/sam3.h index 204e487..2f5c374 100644 --- a/sam3.h +++ b/sam3.h @@ -101,38 +101,6 @@ struct sam3_params { int encode_img_size = 0; // 0 = model default; override input resolution }; -struct sam3_tensor_info { - int64_t ne[4] = {0, 0, 0, 0}; - uint64_t nb[4] = {0, 0, 0, 0}; - int type = 0; - int op = 0; - bool is_contiguous = false; -}; - -enum sam3_vit_block_stage { - SAM3_VIT_BLOCK_STAGE_NORM1 = 0, - SAM3_VIT_BLOCK_STAGE_WINDOW_PART, - SAM3_VIT_BLOCK_STAGE_QKV_PROJ, - SAM3_VIT_BLOCK_STAGE_ATTN_CORE, - SAM3_VIT_BLOCK_STAGE_ATTN_PROJ, - SAM3_VIT_BLOCK_STAGE_WINDOW_UNPART, - SAM3_VIT_BLOCK_STAGE_NORM2, - SAM3_VIT_BLOCK_STAGE_MLP_FC1, - SAM3_VIT_BLOCK_STAGE_MLP_GELU, - SAM3_VIT_BLOCK_STAGE_MLP_FC2, - SAM3_VIT_BLOCK_STAGE_MLP, -}; - -enum sam3_vit_prefix_stage { - SAM3_VIT_PREFIX_STAGE_PATCH_EMBED = 0, - SAM3_VIT_PREFIX_STAGE_PATCH_IM2COL, - SAM3_VIT_PREFIX_STAGE_PATCH_MULMAT_RAW, - SAM3_VIT_PREFIX_STAGE_PATCH_MULMAT, - SAM3_VIT_PREFIX_STAGE_POS_ADD, - SAM3_VIT_PREFIX_STAGE_LN_PRE_NORM, - SAM3_VIT_PREFIX_STAGE_LN_PRE, -}; - struct sam3_pcs_params { std::string text_prompt; std::vector pos_exemplars; @@ -308,193 +276,6 @@ bool sam3_save_mask(const sam3_mask & mask, const std::string & path) sam3_image sam3_decode_video_frame(const std::string & video_path, int frame_index); sam3_video_info sam3_get_video_info(const std::string & video_path); -/***************************************************************************** -** Test and Debug API -** -** Standalone tokenizer, intermediate tensor dumps, and debug utilities. -** These functions are intended for testing and development only. -*****************************************************************************/ - -bool sam3_test_load_tokenizer(const std::string & model_path); -std::vector sam3_test_tokenize(const std::string & text); - -/* -** Run the text encoder on fixed token IDs and dump standard intermediate -** tensors to /.{bin,shape}. -*/ -bool sam3_test_dump_text_encoder(const sam3_model & model, - const std::vector & token_ids, - const std::string & output_dir, - int n_threads = 4); - -/* -** Run the full phase 5 detector path (fusion encoder + DETR decoder + -** dot-product scoring + segmentation head) on an already-encoded image -** and dump intermediate tensors. -*/ -bool sam3_test_dump_phase5(const sam3_model & model, - const sam3_state & state, - const std::vector & token_ids, - const std::string & output_dir, - int n_threads = 4); - -/* -** Run the phase 5 detector from pre-dumped inputs instead of re-running -** the image/text encoders. Isolates detector numerics from earlier phases. -*/ -bool sam3_test_dump_phase5_from_ref_inputs(const sam3_model & model, - const std::vector & token_ids, - const std::string & prephase_ref_dir, - const std::string & phase5_ref_dir, - const std::string & output_dir, - int n_threads = 4); - -/* -** Run the phase 6 prompt encoder + SAM decoder on an already-encoded -** tracker image state and dump intermediate tensors. -*/ -bool sam3_test_dump_phase6(const sam3_model & model, - const sam3_state & state, - const sam3_pvs_params & params, - const std::string & output_dir, - int n_threads = 4); - -/* -** Run the phase 6 prompt encoder + SAM decoder from pre-dumped phase 3 -** tracker features. Isolates phase 6 numerics from earlier phases. -*/ -bool sam3_test_dump_phase6_from_ref_inputs(const sam3_model & model, - const std::string & prephase_ref_dir, - const sam3_pvs_params & params, - const std::string & output_dir, - int n_threads = 4); - -/* -** Run the phase 7 tracker subgraph from pre-dumped case inputs and dump -** intermediate tensors. Case directory produced by dump_phase7_reference.py. -*/ -bool sam3_test_dump_phase7_from_ref_inputs(const sam3_model & model, - const std::string & case_ref_dir, - const std::string & output_dir, - int n_threads = 4); - -/* -** Run the geometry encoder from pre-computed backbone features and dump -** intermediate tensors. Tests exemplar box coordinate encoding against -** Python reference. -*/ -bool sam3_test_dump_geom_enc(const sam3_model & model, - const std::string & prephase_ref_dir, - const sam3_pcs_params & params, - const std::string & output_dir, - int n_threads = 4); - -/* -** Run ONLY the fusion encoder (6 layers) from pre-dumped inputs (image -** features, pos encoding, prompt tokens, attn bias). Dumps per-layer -** outputs for isolated fenc debugging. -*/ -bool sam3_test_fenc_only(const sam3_model & model, - const std::string & ref_dir, - const std::string & output_dir, - int n_threads = 4); - -/* -** ── Debug ──────────────────────────────────────────────────────────────── -*/ - -/* Dump a named state tensor to a binary file for verification. */ -bool sam3_dump_state_tensor(const sam3_state & state, - const std::string & tensor_name, - const std::string & output_path); - -/* Query metadata for a named state tensor without dumping its payload. */ -bool sam3_get_state_tensor_info(const sam3_state & state, - const std::string & tensor_name, - sam3_tensor_info & info); - -/* Dump a named model tensor (weights/constants) to a binary file. */ -bool sam3_dump_model_tensor(const sam3_model & model, - const std::string & tensor_name, - const std::string & output_path); - -/* Query metadata for a named model tensor. */ -bool sam3_get_model_tensor_info(const sam3_model & model, - const std::string & tensor_name, - sam3_tensor_info & info); - -/* -** Encode an image from pre-preprocessed float data (CHW layout, already -** resized and normalized). Bypasses C++ preprocessing so that numerical -** comparisons against the Python reference are not polluted by resize -** implementation differences. -*/ -bool sam3_encode_image_from_preprocessed(sam3_state & state, - const sam3_model & model, - const float * chw_data, - int img_size); - -/* -** Test-only: run ONLY the ViT encoder from preprocessed float data and keep -** only the requested intermediate tensors alive for dumping/comparison. -*/ -bool sam3_encode_vit_from_preprocessed_selective(sam3_state & state, - const sam3_model & model, - const float * chw_data, - int img_size, - const std::vector & output_tensors); - -/* -** Test-only: run the exact ViT prefix up to the tensor entering block 0 -** (patch embed + pos embed + ln_pre). -*/ -bool sam3_test_run_vit_block0_input(const sam3_model & model, - const float * chw_data, - int img_size, - std::vector & output_data, - int64_t output_ne[4], - int n_threads = 4); - -/* -** Test-only: run an exact ViT prefix sub-stage on the real model tensors using -** the model's backend. PATCH_EMBED expects image input [W,H,3,1]. Later stages -** expect feature input [E,W,H,1] in ggml 4D layout. -*/ -bool sam3_test_run_vit_prefix_stage(const sam3_model & model, - sam3_vit_prefix_stage stage, - const float * input_data, - const int64_t input_ne[4], - std::vector & output_data, - int64_t output_ne[4], - int n_threads = 4); -bool sam3_test_run_patch_mulmat_host_ref(const sam3_model & model, - const float * input_data, - const int64_t input_ne[4], - bool use_double_accum, - std::vector & output_data, - int64_t output_ne[4]); -bool sam3_test_run_vit_block_linear_host_ref(const sam3_model & model, - int block_idx, - sam3_vit_block_stage stage, - const float * input_data, - const int64_t input_ne[4], - bool use_double_accum, - std::vector & output_data, - int64_t output_ne[4]); - -/* -** Test-only: run an exact ViT block sub-stage on the real model tensors using -** the model's backend. The input is always provided as F32 in ggml 4D layout. -*/ -bool sam3_test_run_vit_block_stage(const sam3_model & model, - int block_idx, - sam3_vit_block_stage stage, - const float * input_data, - const int64_t input_ne[4], - std::vector & output_data, - int64_t output_ne[4], - int n_threads = 4); - /* ** ── Profiling ─────────────────────────────────────────────────────────── */ diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt deleted file mode 100644 index 8e27787..0000000 --- a/tests/CMakeLists.txt +++ /dev/null @@ -1,129 +0,0 @@ -add_executable(test_load test_load.cpp) -target_link_libraries(test_load PRIVATE sam3) - -add_executable(test_tokenizer test_tokenizer.cpp) -target_link_libraries(test_tokenizer PRIVATE sam3) - -add_executable(test_vit test_vit.cpp) -target_link_libraries(test_vit PRIVATE sam3) - -add_executable(test_phase5 test_phase5.cpp) -target_link_libraries(test_phase5 PRIVATE sam3) - -add_executable(test_phase6 test_phase6.cpp) -target_link_libraries(test_phase6 PRIVATE sam3) - -add_executable(test_metal_phase6 test_metal_phase6.cpp) -target_link_libraries(test_metal_phase6 PRIVATE sam3) - -add_executable(test_phase3 test_phase3.cpp) -target_link_libraries(test_phase3 PRIVATE sam3) - -add_executable(test_phase4 test_phase4.cpp) -target_link_libraries(test_phase4 PRIVATE sam3) - -add_executable(test_phase7 test_phase7.cpp) -target_link_libraries(test_phase7 PRIVATE sam3) - -add_executable(test_debug_encoder test_debug_encoder.cpp) -target_link_libraries(test_debug_encoder PRIVATE sam3) - -add_executable(test_preprocess_match test_preprocess_match.cpp) -target_link_libraries(test_preprocess_match PRIVATE sam3) - -add_executable(test_jpeg_decode test_jpeg_decode.cpp) -target_link_libraries(test_jpeg_decode PRIVATE sam3) - -add_executable(test_geom_enc test_geom_enc.cpp) -target_link_libraries(test_geom_enc PRIVATE sam3) - -add_executable(test_debug_fenc test_debug_fenc.cpp) -target_link_libraries(test_debug_fenc PRIVATE sam3) - -add_executable(test_text_enc_dump test_text_enc_dump.cpp) -target_link_libraries(test_text_enc_dump PRIVATE sam3) - -add_executable(test_e2e_pvs test_e2e_pvs.cpp) -target_link_libraries(test_e2e_pvs PRIVATE sam3) - -add_executable(test_e2e_pcs test_e2e_pcs.cpp) -target_link_libraries(test_e2e_pcs PRIVATE sam3) - -add_executable(test_text_enc_cat test_text_enc_cat.cpp) -target_link_libraries(test_text_enc_cat PRIVATE sam3) - -add_executable(test_llama_pcs test_llama_pcs.cpp) -target_link_libraries(test_llama_pcs PRIVATE sam3) - -add_executable(test_pvs_point test_pvs_point.cpp) -target_link_libraries(test_pvs_point PRIVATE sam3) - -add_executable(test_video_2frame test_video_2frame.cpp) -target_link_libraries(test_video_2frame PRIVATE sam3) - -add_executable(test_video_text test_video_text.cpp) -target_link_libraries(test_video_text PRIVATE sam3) - -add_executable(test_metal_win test_metal_win.cpp) -target_link_libraries(test_metal_win PRIVATE sam3) - -add_executable(test_metal_vit test_metal_vit.cpp) -target_link_libraries(test_metal_vit PRIVATE sam3) - -add_executable(test_metal_encoder_checkpoints test_metal_encoder_checkpoints.cpp) -target_link_libraries(test_metal_encoder_checkpoints PRIVATE sam3) - -add_executable(test_metal_flash_attn_ext_block15_compare test_metal_flash_attn_ext_block15_compare.cpp) -target_link_libraries(test_metal_flash_attn_ext_block15_compare PRIVATE sam3) - -add_executable(test_metal_vit_upstream_compare test_metal_vit_upstream_compare.cpp) -target_link_libraries(test_metal_vit_upstream_compare PRIVATE sam3) - -add_executable(test_metal_vit_prefix_compare test_metal_vit_prefix_compare.cpp) -target_link_libraries(test_metal_vit_prefix_compare PRIVATE sam3) - -add_executable(test_metal_block0_stage_compare test_metal_block0_stage_compare.cpp) -target_link_libraries(test_metal_block0_stage_compare PRIVATE sam3) - -add_executable(test_metal_block_stage_compare test_metal_block_stage_compare.cpp) -target_link_libraries(test_metal_block_stage_compare PRIVATE sam3) - -add_executable(test_metal_python_ref_compare test_metal_python_ref_compare.cpp) -target_link_libraries(test_metal_python_ref_compare PRIVATE sam3) - -add_executable(test_metal_conv_dw test_metal_conv_dw.cpp) -target_link_libraries(test_metal_conv_dw PRIVATE sam3) - -add_executable(test_metal_profile test_metal_profile.cpp) -target_link_libraries(test_metal_profile PRIVATE sam3) - -add_executable(test_metal_overhead test_metal_overhead.cpp) -target_link_libraries(test_metal_overhead PRIVATE sam3) - -add_executable(test_v_cont_removal test_v_cont_removal.cpp) -target_link_libraries(test_v_cont_removal PRIVATE sam3) - -add_executable(test_metal_real_profile test_metal_real_profile.cpp) -target_link_libraries(test_metal_real_profile PRIVATE sam3) - -add_executable(test_metal_conv_transpose_2d_stress test_metal_conv_transpose_2d_stress.cpp) -target_link_libraries(test_metal_conv_transpose_2d_stress PRIVATE sam3) - -add_executable(test_metal_conv_transpose_2d test_metal_conv_transpose_2d.cpp) -target_link_libraries(test_metal_conv_transpose_2d PRIVATE sam3) - -add_executable(test_visual_only test_visual_only.cpp) -target_link_libraries(test_visual_only PRIVATE sam3) - -add_executable(test_visual_only_compare test_visual_only_compare.cpp) -target_link_libraries(test_visual_only_compare PRIVATE sam3) - -add_executable(test_sam2_backbone_compare test_sam2_backbone_compare.cpp) -target_link_libraries(test_sam2_backbone_compare PRIVATE sam3) - -add_executable(test_sam2_pe_cache_size test_sam2_pe_cache_size.cpp) -target_link_libraries(test_sam2_pe_cache_size PRIVATE sam3) - -add_executable(test_sam2_video_debug test_sam2_video_debug.cpp) -target_link_libraries(test_sam2_video_debug PRIVATE sam3) - diff --git a/tests/analyze_divergence.py b/tests/analyze_divergence.py deleted file mode 100644 index 1b57b9a..0000000 --- a/tests/analyze_divergence.py +++ /dev/null @@ -1,114 +0,0 @@ -#!/usr/bin/env python3 -"""Analyze divergence between C++ and Python ViT outputs. - -Focuses on understanding the error distribution and whether it's -just f32 precision on large-magnitude values. -""" -import numpy as np -import os, sys - -def load_tensor(path): - with open(path + ".shape") as f: - shape = [int(x) for x in f.read().strip().split(",") if x] - data = np.fromfile(path + ".bin", dtype=np.float32) - return data, shape - -ref_dir = sys.argv[1] if len(sys.argv) > 1 else "tests/ref_phase3" -cpp_dir = sys.argv[2] if len(sys.argv) > 2 else "tests/cpp_out" - -print("═══ Divergence Analysis ═══\n") - -# Analyze per-block error progression -print("Block-by-block error progression:") -print(f"{'Block':>6} {'MaxErr':>12} {'MeanErr':>12} {'MaxVal':>12} {'RelErr':>12} {'Cosine':>12}") -print("-" * 78) - -for i in range(32): - ref_path = f"{ref_dir}/block_{i}_out" - cpp_path = f"{cpp_dir}/dbg_block_{i}_out" - if not os.path.exists(ref_path + ".bin") or not os.path.exists(cpp_path + ".bin"): - continue - - ref, ref_shape = load_tensor(ref_path) - cpp, cpp_shape = load_tensor(cpp_path) - - # ggml [E, W, H] → NHWC [H, W, E] is the same flat layout for comparison - diff = np.abs(ref - cpp) - max_err = diff.max() - mean_err = diff.mean() - max_val = max(np.abs(ref).max(), np.abs(cpp).max()) - rel_err = max_err / max_val if max_val > 0 else 0 - cos = np.dot(ref, cpp) / (np.linalg.norm(ref) * np.linalg.norm(cpp) + 1e-12) - - print(f"{i:>6} {max_err:>12.6e} {mean_err:>12.6e} {max_val:>12.4f} {rel_err:>12.6e} {cos:>12.10f}") - -# Detailed analysis of the worst element -print("\n═══ Worst Element Analysis ═══\n") -for block in [14, 15, 23, 31]: - ref_path = f"{ref_dir}/block_{block}_out" - cpp_path = f"{cpp_dir}/dbg_block_{block}_out" - if block == 31: - cpp_path = f"{cpp_dir}/vit_output" - if not os.path.exists(ref_path + ".bin") or not os.path.exists(cpp_path + ".bin"): - continue - - ref, ref_shape = load_tensor(ref_path) - cpp, cpp_shape = load_tensor(cpp_path) - - # For block 31 / vit_output, need NCHW→NHWC transpose - if block == 31 and len(cpp_shape) == 3: - # cpp is [C, W, H] ggml → ref is [1, H, W, C] NHWC - C, W, H = cpp_shape - cpp_3d = cpp.reshape(H, W, C) - # ref is [1, H, W, C] NHWC - ref_3d = ref.reshape(ref_shape[1], ref_shape[2], ref_shape[3]) - diff = np.abs(ref_3d.flatten() - cpp_3d.flatten()) - else: - diff = np.abs(ref - cpp) - - worst_idx = np.argmax(diff) - - # Convert flat index to spatial coordinates - # ref is [1, H=72, W=72, E=1024] - if len(ref_shape) == 4: - _, H, W, E = ref_shape - else: - H, W, E = ref_shape[-3], ref_shape[-2], ref_shape[-1] - - total = H * W * E - h = worst_idx // (W * E) - w = (worst_idx % (W * E)) // E - e = worst_idx % E - - print(f"Block {block}: worst at (h={h}, w={w}, e={e})") - print(f" ref={ref[worst_idx]:.6e}, cpp={cpp[worst_idx]:.6e}, diff={diff[worst_idx]:.6e}") - print(f" magnitude={abs(ref[worst_idx]):.4f}, rel_err={diff[worst_idx]/abs(ref[worst_idx]):.6e}") - - # Check: how many elements have error > 1e-3? - n_above_1e3 = (diff > 1e-3).sum() - n_above_1e2 = (diff > 1e-2).sum() - print(f" elements with |err| > 1e-3: {n_above_1e3} / {len(diff)} ({100*n_above_1e3/len(diff):.4f}%)") - print(f" elements with |err| > 1e-2: {n_above_1e2} / {len(diff)} ({100*n_above_1e2/len(diff):.4f}%)") - print() - -# Compare block 15 intermediates -print("\n═══ Block 15 Intermediate Comparison ═══\n") -intermediates = [ - ("after_norm1", "norm1"), - ("q_pre_rope", "Q pre-RoPE"), - ("k_pre_rope", "K pre-RoPE"), - ("q_post_rope", "Q post-RoPE"), - ("k_post_rope", "K post-RoPE"), - ("v", "V"), - ("attn_out", "attention output"), - ("proj_out", "projection output"), - ("after_attn_residual", "after attn+residual"), -] - -for suffix, label in intermediates: - ref_path = f"{ref_dir}/block_15_{suffix}" - if not os.path.exists(ref_path + ".bin"): - continue - ref, ref_shape = load_tensor(ref_path) - max_val = np.abs(ref).max() - print(f" block_15_{suffix:25s} shape={ref_shape} range=[{ref.min():.4f}, {ref.max():.4f}] max_abs={max_val:.4f}") diff --git a/tests/analyze_f16_divergence.py b/tests/analyze_f16_divergence.py deleted file mode 100644 index 84d7a87..0000000 --- a/tests/analyze_f16_divergence.py +++ /dev/null @@ -1,96 +0,0 @@ -#!/usr/bin/env python3 -"""Analyze where f16 precision breaks down in the ViT blocks.""" -import numpy as np -import os - -ref_dir = "tests/ref_phase3" -cpp_dir = "tests/cpp_out_f16" - -print("═══ f16 C++ vs f32 Python — Block-by-Block Error Progression ═══\n") -print(f"{'Block':>6} {'MaxErr':>12} {'MeanErr':>12} {'MaxVal':>12} {'RelErr':>12} {'Cosine':>14} {'WorstIdx':>10}") -print("-" * 90) - -for i in range(32): - ref_path = f"{ref_dir}/block_{i}_out" - cpp_path = f"{cpp_dir}/dbg_block_{i}_out" - if not os.path.exists(ref_path + ".bin") or not os.path.exists(cpp_path + ".bin"): - continue - - ref = np.fromfile(ref_path + ".bin", dtype=np.float32) - cpp = np.fromfile(cpp_path + ".bin", dtype=np.float32) - - diff = np.abs(ref - cpp) - max_err = diff.max() - mean_err = diff.mean() - max_val = max(np.abs(ref).max(), np.abs(cpp).max()) - rel_err = max_err / max_val if max_val > 0 else 0 - cos = np.dot(ref, cpp) / (np.linalg.norm(ref) * np.linalg.norm(cpp) + 1e-12) - worst = np.argmax(diff) - - flag = " <<<" if i in [10, 11, 15, 16] or max_err > 1.0 else "" - print(f"{i:>6} {max_err:>12.6e} {mean_err:>12.6e} {max_val:>12.4f} {rel_err:>12.6e} {cos:>14.10f} {worst:>10}{flag}") - -# Detailed analysis of the jump between block 14 and 15 -print("\n═══ Detailed Block 15 Analysis (first global block that breaks) ═══\n") - -# Check block 15 intermediates -intermediates = [ - "block_15_after_norm1", - "block_15_q_pre_rope", "block_15_k_pre_rope", - "block_15_q_post_rope", "block_15_k_post_rope", - "block_15_v", - "block_15_attn_out", - "block_15_proj_out", - "block_15_after_attn_residual", -] - -for name in intermediates: - ref_path = f"{ref_dir}/{name}" - if not os.path.exists(ref_path + ".bin"): - continue - ref = np.fromfile(ref_path + ".bin", dtype=np.float32) - print(f" {name:40s} range=[{ref.min():.4f}, {ref.max():.4f}] " - f"max_abs={np.abs(ref).max():.4f}") - -# Check the worst element across blocks -print("\n═══ Tracking the Worst Element ═══\n") -for i in [10, 11, 12, 13, 14, 15, 16, 20, 23, 30]: - ref_path = f"{ref_dir}/block_{i}_out" - cpp_path = f"{cpp_dir}/dbg_block_{i}_out" - if not os.path.exists(ref_path + ".bin") or not os.path.exists(cpp_path + ".bin"): - continue - - ref = np.fromfile(ref_path + ".bin", dtype=np.float32) - cpp = np.fromfile(cpp_path + ".bin", dtype=np.float32) - diff = np.abs(ref - cpp) - worst = np.argmax(diff) - - # Convert flat index to spatial coords (ref is [1, 72, 72, 1024] NHWC) - H, W, E = 72, 72, 1024 - h = worst // (W * E) - w = (worst % (W * E)) // E - e = worst % E - - print(f" Block {i:2d}: worst at (h={h:2d}, w={w:2d}, e={e:4d}) " - f"ref={ref[worst]:>12.4f} cpp={cpp[worst]:>12.4f} " - f"diff={diff[worst]:>10.4f} rel={diff[worst]/abs(ref[worst]) if ref[worst] != 0 else 0:.6f}") - -# Check: how many elements have error > 1.0? -print("\n═══ Error Distribution by Block ═══\n") -print(f"{'Block':>6} {'> 0.01':>10} {'> 0.1':>10} {'> 1.0':>10} {'> 10.0':>10}") -for i in range(32): - ref_path = f"{ref_dir}/block_{i}_out" - cpp_path = f"{cpp_dir}/dbg_block_{i}_out" - if not os.path.exists(ref_path + ".bin") or not os.path.exists(cpp_path + ".bin"): - continue - - ref = np.fromfile(ref_path + ".bin", dtype=np.float32) - cpp = np.fromfile(cpp_path + ".bin", dtype=np.float32) - diff = np.abs(ref - cpp) - - n01 = (diff > 0.01).sum() - n1 = (diff > 0.1).sum() - n10 = (diff > 1.0).sum() - n100 = (diff > 10.0).sum() - flag = " <<<" if n10 > 0 else "" - print(f"{i:>6} {n01:>10} {n1:>10} {n10:>10} {n100:>10}{flag}") diff --git a/tests/cat.jpg b/tests/cat.jpg deleted file mode 100644 index edecb9c..0000000 Binary files a/tests/cat.jpg and /dev/null differ diff --git a/tests/check_f16_weight_error.py b/tests/check_f16_weight_error.py deleted file mode 100644 index 0af1f04..0000000 --- a/tests/check_f16_weight_error.py +++ /dev/null @@ -1,180 +0,0 @@ -#!/usr/bin/env python3 -"""Check f16 quantization error on ViT weights, focusing on dimension 679. - -Also: run the same ViT forward pass in Python with f16 weights to see if -PyTorch shows the same amplification. This tells us if the issue is inherent -to f16 or a C++ bug. -""" -import numpy as np -import torch -import os - -# Load checkpoint -print("Loading checkpoint...") -ckpt = torch.load("raw_weights/sam3.pt", map_location="cpu", weights_only=False) -if "model" in ckpt: - ckpt = ckpt["model"] - -vit_prefix = "detector.backbone.vision_backbone.trunk." - -# Check f16 quantization error on specific block weights -print("\n═══ F16 Quantization Error on ViT Weights ═══\n") -print(f"{'Weight':50s} {'MaxRelErr':>12} {'MaxAbsErr':>12} {'ErrAt679':>12} {'MaxVal':>12}") -print("-" * 100) - -for blk in [0, 7, 14, 15, 16, 31]: - for suffix in ["attn.qkv.weight", "attn.qkv.bias", "attn.proj.weight", - "mlp.fc1.weight", "mlp.fc2.weight"]: - key = f"{vit_prefix}blocks.{blk}.{suffix}" - if key not in ckpt: - continue - w = ckpt[key].float() - w_f16 = w.half().float() - err = (w - w_f16).abs() - max_abs = err.max().item() - max_rel = (err / (w.abs() + 1e-12)).max().item() - max_val = w.abs().max().item() - - # Check dimension 679 specifically - if w.ndim == 2 and w.shape[1] > 679: - err_679 = err[:, 679].max().item() - elif w.ndim == 2 and w.shape[0] > 679: - err_679 = err[679, :].max().item() - elif w.ndim == 1 and w.shape[0] > 679: - err_679 = err[679].item() - else: - err_679 = 0.0 - - name = f"blk{blk}.{suffix}" - print(f" {name:48s} {max_rel:>12.6e} {max_abs:>12.6e} {err_679:>12.6e} {max_val:>12.4f}") - -# ═══ Compare: Python f32 vs Python f16 ViT forward ═══ -print("\n═══ Python f32 vs f16 ViT Forward Pass ═══") -print("(This tells us if the error is inherent to f16 or a C++ bug)\n") - -# Load the preprocessed image -ref_dir = "tests/ref_phase3" -img_data = np.fromfile(f"{ref_dir}/preprocessed.bin", dtype=np.float32) -img = torch.from_numpy(img_data).reshape(1, 3, 1008, 1008) - -# Extract ViT weights -vit_sd = {k[len(vit_prefix):]: v for k, v in ckpt.items() if k.startswith(vit_prefix)} - -# Run ViT forward pass with f32 and f16 weights, compare outputs -# We'll just check the QKV output at block 15 and the block output - -from torchvision.transforms import v2 -import torch.nn.functional as F -import math - -def get_abs_pos(abs_pos, has_cls_token, hw, tiling=False): - h, w = hw - if has_cls_token: - abs_pos = abs_pos[:, 1:] - xy_num = abs_pos.shape[1] - size = int(math.sqrt(xy_num)) - if size != h or size != w: - new_abs_pos = abs_pos.reshape(1, size, size, -1).permute(0, 3, 1, 2) - if tiling: - new_abs_pos = new_abs_pos.tile( - [1, 1] + [x // y + 1 for x, y in zip((h, w), new_abs_pos.shape[2:])] - )[:, :, :h, :w] - return new_abs_pos.permute(0, 2, 3, 1) - return abs_pos.reshape(1, h, w, -1) - -def apply_rotary_enc(xq, xk, freqs_cis): - xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) - xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2)) - freqs_cis = freqs_cis.view(*[d if i >= xq_.ndim - 2 else 1 for i, d in enumerate(xq_.shape)]) - xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3) - xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3) - return xq_out.type_as(xq), xk_out.type_as(xk) - -def window_partition(x, ws): - B, H, W, C = x.shape - x = x.view(B, H // ws, ws, W // ws, ws, C) - return x.permute(0, 1, 3, 2, 4, 5).reshape(-1, ws, ws, C), (H, W) - -def window_unpartition(windows, ws, pad_hw, hw): - Hp, Wp = pad_hw - H, W = hw - B = windows.shape[0] // (Hp * Wp // ws // ws) - x = windows.reshape(B, Hp // ws, Wp // ws, ws, ws, -1) - return x.permute(0, 1, 3, 2, 4, 5).reshape(B, Hp, Wp, -1)[:, :H, :W, :] - -def run_vit(vit_sd, img, use_f16_weights=False): - """Run ViT forward pass. If use_f16_weights, quantize all 2D weights to f16.""" - E, NH, HD, WS = 1024, 16, 64, 24 - global_blocks = {7, 15, 23, 31} - - # Optionally quantize weights to f16 - sd = {} - for k, v in vit_sd.items(): - if use_f16_weights and v.ndim >= 2: - sd[k] = v.half().float() # f32 → f16 → f32 roundtrip - else: - sd[k] = v.float() - - with torch.no_grad(): - patch_w = sd["patch_embed.proj.weight"] - x = F.conv2d(img, patch_w, stride=14).permute(0, 2, 3, 1) - pos = get_abs_pos(sd["pos_embed"], True, (72, 72), tiling=True) - x = x + pos - x = F.layer_norm(x, [E], sd["ln_pre.weight"], sd["ln_pre.bias"]) - - block_outs = {} - for blk_idx in range(32): - prefix = f"blocks.{blk_idx}." - is_global = blk_idx in global_blocks - ws = 0 if is_global else WS - shortcut = x - xn = F.layer_norm(x, [E], sd[prefix + "norm1.weight"], sd[prefix + "norm1.bias"]) - if ws > 0: - H, W = xn.shape[1], xn.shape[2] - xn, pad_hw = window_partition(xn, ws) - B_cur, Hc, Wc, _ = xn.shape - L = Hc * Wc - qkv = F.linear(xn, sd[prefix + "attn.qkv.weight"], sd[prefix + "attn.qkv.bias"]) - qkv = qkv.reshape(B_cur, L, 3, NH, HD) - q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0) - freqs_key = prefix + "attn.freqs_cis" - if freqs_key in sd: - q, k = apply_rotary_enc(q, k, sd[freqs_key]) - attn_out = F.scaled_dot_product_attention(q, k, v) - attn_out = attn_out.view(B_cur, NH, Hc, Wc, HD).permute(0, 2, 3, 1, 4).reshape(B_cur, Hc, Wc, E) - attn_out = F.linear(attn_out, sd[prefix + "attn.proj.weight"], sd[prefix + "attn.proj.bias"]) - if ws > 0: - attn_out = window_unpartition(attn_out, ws, pad_hw, (H, W)) - x = shortcut + attn_out - shortcut = x - xn = F.layer_norm(x, [E], sd[prefix + "norm2.weight"], sd[prefix + "norm2.bias"]) - h = F.gelu(F.linear(xn, sd[prefix + "mlp.fc1.weight"], sd[prefix + "mlp.fc1.bias"])) - h = F.linear(h, sd[prefix + "mlp.fc2.weight"], sd[prefix + "mlp.fc2.bias"]) - x = shortcut + h - block_outs[blk_idx] = x.clone() - - return block_outs - -print("Running f32 forward pass...") -outs_f32 = run_vit(vit_sd, img, use_f16_weights=False) -print("Running f16-weights forward pass...") -outs_f16 = run_vit(vit_sd, img, use_f16_weights=True) - -print(f"\n{'Block':>6} {'MaxErr':>12} {'MeanErr':>12} {'Cosine':>14}") -print("-" * 50) -for blk in sorted(outs_f32.keys()): - a = outs_f32[blk].flatten().numpy() - b = outs_f16[blk].flatten().numpy() - diff = np.abs(a - b) - cos = np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)) - print(f"{blk:>6} {diff.max():>12.6e} {diff.mean():>12.6e} {cos:>14.10f}") - -# Check the worst element at dimension 679 -print(f"\n═══ Worst Element (dim 679) in Python f16 ═══") -for blk in [14, 15, 16, 23, 30]: - a = outs_f32[blk][0].numpy() # [72, 72, 1024] - b = outs_f16[blk][0].numpy() - diff = np.abs(a - b) - worst = np.unravel_index(np.argmax(diff), diff.shape) - print(f" Block {blk:2d}: worst at (h={worst[0]:2d}, w={worst[1]:2d}, e={worst[2]:4d}) " - f"f32={a[worst]:>10.4f} f16={b[worst]:>10.4f} diff={diff[worst]:>8.4f}") diff --git a/tests/check_preprocessing.py b/tests/check_preprocessing.py deleted file mode 100644 index 1875724..0000000 --- a/tests/check_preprocessing.py +++ /dev/null @@ -1,73 +0,0 @@ -#!/usr/bin/env python3 -"""Check if the preprocessing difference explains the ViT output difference.""" -import numpy as np -import sys - -def load_tensor(path): - with open(path + ".shape") as f: - shape = [int(x) for x in f.read().strip().split(",")] - with open(path + ".bin", "rb") as f: - data = np.frombuffer(f.read(), dtype=np.float32).copy() - return data.reshape(shape) - -ref_dir = sys.argv[1] if len(sys.argv) > 1 else "tests/ref_phase3" - -# Check if C++ preprocessed image is available -try: - cpp_prep = load_tensor(f"{ref_dir}/cpp_out_phase3/preprocessed") - py_prep = load_tensor(f"{ref_dir}/preprocessed") - diff = np.abs(cpp_prep - py_prep) - print(f"Preprocessing comparison:") - print(f" C++ shape: {cpp_prep.shape}, range=[{cpp_prep.min():.6f}, {cpp_prep.max():.6f}]") - print(f" Py shape: {py_prep.shape}, range=[{py_prep.min():.6f}, {py_prep.max():.6f}]") - print(f" max_diff: {diff.max():.6f}") - print(f" mean_diff: {diff.mean():.6f}") -except: - print("C++ preprocessed tensor not available. Checking against known resize differences.") - print() - print("The C++ bilinear resize operates on uint8 values, then normalizes.") - print("Python v2.Resize uses PIL/pillow internally which may use different algorithms.") - print() - - # Check how different the two resize methods could be - from PIL import Image - import torch - from torchvision.transforms import v2 - - img = Image.open("tests/test_random.jpg").convert("RGB") - - # Python's method - transform = v2.Compose([ - v2.ToDtype(torch.uint8, scale=True), - v2.Resize(size=(1008, 1008)), - v2.ToDtype(torch.float32, scale=True), - v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), - ]) - img_tensor = v2.functional.to_image(img) - py_result = transform(img_tensor).numpy() - - # Manual bilinear (approximating C++) - import numpy as np - from PIL import Image - arr = np.array(img) - h, w = arr.shape[:2] - - # PIL resize with BILINEAR - img_pil = Image.fromarray(arr).resize((1008, 1008), Image.BILINEAR) - pil_arr = np.array(img_pil).astype(np.float32) / 255.0 - pil_normalized = (pil_arr - 0.5) / 0.5 # [H, W, C] - pil_chw = np.transpose(pil_normalized, (2, 0, 1)) # [C, H, W] - - diff = np.abs(pil_chw - py_result) - print(f"PIL BILINEAR vs torchvision Resize:") - print(f" max_diff: {diff.max():.6f}") - print(f" mean_diff: {diff.mean():.6f}") - print(f" This is approximately 1/127.5 = {1/127.5:.6f} per uint8 step") - - # Check v2.Resize method - print(f"\nPython transform output range: [{py_result.min():.4f}, {py_result.max():.4f}]") - print(f"PIL resize output range: [{pil_chw.min():.4f}, {pil_chw.max():.4f}]") - -print("\nKey insight: bilinear resize differences produce ~0.01 pixel-level errors") -print("After 32 transformer blocks, these can amplify significantly.") -print("This is EXPECTED behavior, not a bug.") diff --git a/tests/check_rounding.py b/tests/check_rounding.py deleted file mode 100644 index 8864b41..0000000 --- a/tests/check_rounding.py +++ /dev/null @@ -1,31 +0,0 @@ -#!/usr/bin/env python3 -"""Check torchvision's uint8 rounding behavior for bilinear resize.""" -import torch -from torchvision.transforms.functional import resize -from torchvision.transforms import InterpolationMode - -# Create a simple test image where we know the exact bilinear result -# Use a 2x2 image resized to 4x4 - the bilinear result should be exactly deterministic -img = torch.tensor([[[100, 200], [50, 150]]], dtype=torch.uint8) # [1, 2, 2] -print(f"Input: {img}") - -# Resize to 4x4 -out = resize(img, [4, 4], interpolation=InterpolationMode.BILINEAR, antialias=False) -print(f"Output (antialias=False): {out}") - -out_aa = resize(img, [4, 4], interpolation=InterpolationMode.BILINEAR, antialias=True) -print(f"Output (antialias=True): {out_aa}") - -# Check rounding: does torch round or truncate? -# For a pixel value that should be 127.5 after bilinear: -# Truncate → 127, Round → 128 -img2 = torch.tensor([[[0, 255], [0, 255]]], dtype=torch.uint8) -out2 = resize(img2, [2, 3], interpolation=InterpolationMode.BILINEAR, antialias=False) -print(f"\nRounding test input: {img2}") -print(f"Rounding test output: {out2}") -# Middle column should be (0+255)/2 = 127.5 → check if 127 or 128 - -# Also check: does torchvision Resize on uint8 go through float internally? -import torchvision -print(f"\ntorchvision version: {torchvision.__version__}") -print(f"torch version: {torch.__version__}") diff --git a/tests/compare_actual_resize.py b/tests/compare_actual_resize.py deleted file mode 100644 index 06da279..0000000 --- a/tests/compare_actual_resize.py +++ /dev/null @@ -1,51 +0,0 @@ -#!/usr/bin/env python3 -"""Compare ACTUAL C++ resized uint8 (from sam3_preprocess_image dump) -against Python torchvision v2.Resize uint8 output.""" -import numpy as np - -# Load C++ resized uint8 [H, W, 3] -cpp = np.fromfile("tests/ref_phase3/cpp_resized_uint8.bin", dtype=np.uint8).reshape(1008, 1008, 3) - -# Load Python torchvision resized uint8 [H, W, 3] -py = np.fromfile("tests/ref_phase3/resized_uint8.bin", dtype=np.uint8).reshape(1008, 1008, 3) - -print("═══ Actual C++ resize vs Python torchvision resize (uint8) ═══") -print(f" C++ shape: {cpp.shape}, range: [{cpp.min()}, {cpp.max()}]") -print(f" Python shape: {py.shape}, range: [{py.min()}, {py.max()}]") - -diff = np.abs(cpp.astype(np.int16) - py.astype(np.int16)) -print(f"\n max pixel diff: {diff.max()}") -print(f" mean pixel diff: {diff.mean():.6f}") -n_diff = (diff > 0).sum() -print(f" pixels differing: {n_diff} / {diff.size} ({100*n_diff/diff.size:.2f}%)") - -for t in [0, 1, 2, 3, 5, 10]: - n = (diff > t).sum() - if n > 0: - print(f" pixels with diff > {t}: {n}") - -# Find worst locations -if diff.max() > 0: - worst_flat = np.argmax(diff.flatten()) - worst_pos = np.unravel_index(worst_flat, diff.shape) - print(f"\n Worst pixel at (y={worst_pos[0]}, x={worst_pos[1]}, c={worst_pos[2]})") - print(f" C++={cpp[worst_pos]}, Python={py[worst_pos]}, diff={diff[worst_pos]}") - - # Show some worst locations - flat_diff = diff.flatten() - worst_indices = np.argsort(flat_diff)[-10:][::-1] - print(f"\n Top 10 worst positions:") - for idx in worst_indices: - pos = np.unravel_index(idx, diff.shape) - print(f" y={pos[0]:4d} x={pos[1]:4d} c={pos[2]} cpp={cpp[pos]:3d} py={py[pos]:3d} diff={diff[pos]:2d}") - -# After fix: compare normalized float outputs -print("\n═══ Normalized float comparison ═══") -cpp_float = (cpp.astype(np.float32) / 255.0 - 0.5) / 0.5 -py_float = (py.astype(np.float32) / 255.0 - 0.5) / 0.5 -float_diff = np.abs(cpp_float - py_float) -print(f" max float diff: {float_diff.max():.6e}") -print(f" mean float diff: {float_diff.mean():.6e}") -cos = np.dot(cpp_float.flatten(), py_float.flatten()) / ( - np.linalg.norm(cpp_float.flatten()) * np.linalg.norm(py_float.flatten())) -print(f" cosine: {cos:.10f}") diff --git a/tests/compare_cpp_dump_vs_finterp.py b/tests/compare_cpp_dump_vs_finterp.py deleted file mode 100644 index e7afdd1..0000000 --- a/tests/compare_cpp_dump_vs_finterp.py +++ /dev/null @@ -1,62 +0,0 @@ -#!/usr/bin/env python3 -"""Compare the ACTUAL C++ uint8 dump against torch F.interpolate output.""" -import numpy as np -import torch -from PIL import Image -from torchvision.transforms import v2 - -image_path = "tests/test_random.jpg" -img_size = 1008 - -# Load image for torch -img = Image.open(image_path).convert("RGB") -img_tensor = v2.functional.to_image(img) # uint8 CHW [3, 480, 640] - -# ═══ Ground truth: v2.Resize (antialias=True) ═══ -tv_resized = v2.Resize(size=(img_size, img_size))(img_tensor) -tv_hwc = tv_resized.permute(1, 2, 0).numpy() # [H, W, 3] uint8 - -# ═══ F.interpolate (antialias=False) ═══ -float_input = img_tensor.float() -fi_resized = torch.nn.functional.interpolate( - float_input.unsqueeze(0), size=(img_size, img_size), - mode='bilinear', align_corners=False, antialias=False -).squeeze(0) -fi_uint8 = fi_resized.clamp(0, 255).round().to(torch.uint8) -fi_hwc = fi_uint8.permute(1, 2, 0).numpy() - -# ═══ F.interpolate (antialias=True) ═══ -fi_aa_resized = torch.nn.functional.interpolate( - float_input.unsqueeze(0), size=(img_size, img_size), - mode='bilinear', align_corners=False, antialias=True -).squeeze(0) -fi_aa_uint8 = fi_aa_resized.clamp(0, 255).round().to(torch.uint8) -fi_aa_hwc = fi_aa_uint8.permute(1, 2, 0).numpy() - -# ═══ Load ACTUAL C++ dump ═══ -cpp_hwc = np.fromfile("tests/ref_phase3/cpp_resized_uint8.bin", dtype=np.uint8).reshape(img_size, img_size, 3) - -print("═══ Comparisons ═══\n") - -for name, ref in [("v2.Resize", tv_hwc), - ("F.interp(aa=False)", fi_hwc), - ("F.interp(aa=True)", fi_aa_hwc)]: - diff = np.abs(cpp_hwc.astype(np.int16) - ref.astype(np.int16)) - n_diff = (diff > 0).sum() - print(f" C++ vs {name:22s}: max={diff.max():2d} n_diff={n_diff:7d} ({100*n_diff/diff.size:.2f}%)" - f" >1: {(diff>1).sum():6d} >2: {(diff>2).sum():4d}") - -# Also compare the torch methods against each other -for name, ref, comp in [("v2.Resize vs F(aa=F)", tv_hwc, fi_hwc), - ("v2.Resize vs F(aa=T)", tv_hwc, fi_aa_hwc), - ("F(aa=F) vs F(aa=T)", fi_hwc, fi_aa_hwc)]: - diff = np.abs(ref.astype(np.int16) - comp.astype(np.int16)) - n_diff = (diff > 0).sum() - print(f" {name:22s}: max={diff.max():2d} n_diff={n_diff:7d} ({100*n_diff/diff.size:.2f}%)") - -# ═══ Save the F.interpolate(aa=True) + round output for C++ to compare against ═══ -with open("tests/ref_phase3/resized_uint8_finterp_aa.bin", "wb") as f: - f.write(fi_aa_hwc.tobytes()) -with open("tests/ref_phase3/resized_uint8_finterp_aa.shape", "w") as f: - f.write(f"{img_size},{img_size},3") -print(f"\nSaved F.interp(aa=True) uint8 output") diff --git a/tests/compare_ddec_tensors.py b/tests/compare_ddec_tensors.py deleted file mode 100644 index cca04c0..0000000 --- a/tests/compare_ddec_tensors.py +++ /dev/null @@ -1,218 +0,0 @@ -#!/usr/bin/env python3 -"""Comprehensive comparison of DETR decoder tensors: Python ref vs C++ output. - -Computes all 6 metrics: mae, max, rel, cos, p95, p99. -Handles layout differences between PyTorch (seq-first/batch-first) and ggml. -""" -import os -import sys -import numpy as np - -REF_DIR = os.path.join(os.path.dirname(__file__), "ref_phase5") -CPP_DIR = os.path.join(REF_DIR, "cpp_out_phase5") - - -def load_tensor(path): - """Load a .bin + .shape tensor pair.""" - shape_file = path + ".shape" - bin_file = path + ".bin" - if not os.path.exists(shape_file) or not os.path.exists(bin_file): - return None, None - with open(shape_file) as f: - shape = tuple(int(x) for x in f.read().strip().split(",") if x) - data = np.fromfile(bin_file, dtype=np.float32) - return data, shape - - -def compare(name, ref_path, cpp_path, atol=1e-4): - """Compare two tensors and return detailed metrics.""" - ref_data, ref_shape = load_tensor(ref_path) - cpp_data, cpp_shape = load_tensor(cpp_path) - - if ref_data is None: - return None, f"SKIP (missing ref: {ref_path})" - if cpp_data is None: - return None, f"SKIP (missing cpp: {cpp_path})" - - if ref_data.size != cpp_data.size: - return None, f"FAIL NUMEL MISMATCH ref={ref_data.size} ({ref_shape}) cpp={cpp_data.size} ({cpp_shape})" - - diff = np.abs(ref_data - cpp_data) - eps = 1e-8 - - mae = float(diff.mean()) - max_d = float(diff.max()) - rel_err = diff / (np.abs(ref_data) + eps) - mean_rel = float(rel_err.mean()) - cos_num = float(np.dot(ref_data, cpp_data)) - cos_den = float(np.linalg.norm(ref_data) * np.linalg.norm(cpp_data) + 1e-12) - cos = cos_num / cos_den - p95 = float(np.percentile(diff, 95)) - p99 = float(np.percentile(diff, 99)) - - worst_idx = int(np.argmax(diff)) - status = "PASS" if max_d < atol else "FAIL" - - metrics = { - "mae": mae, - "max": max_d, - "rel": mean_rel, - "cos": cos, - "p95": p95, - "p99": p99, - "worst_idx": worst_idx, - "worst_ref": float(ref_data[worst_idx]), - "worst_cpp": float(cpp_data[worst_idx]), - "n_over_atol": int(np.sum(diff > atol)), - "n_total": int(ref_data.size), - "ref_shape": ref_shape, - "cpp_shape": cpp_shape, - } - return metrics, status - - -def print_result(name, metrics, status, atol): - if metrics is None: - print(f" {status}") - return - - label = f"{status} {name:45s} ref_shape={metrics['ref_shape']} cpp_shape={metrics['cpp_shape']}" - print(f" {label}") - print(f" mae={metrics['mae']:.6e} max={metrics['max']:.6e} " - f"rel={metrics['rel']:.6e} cos={metrics['cos']:.10f} " - f"p95={metrics['p95']:.6e} p99={metrics['p99']:.6e}") - if status == "FAIL": - print(f" worst_idx={metrics['worst_idx']} " - f"ref={metrics['worst_ref']:.6e} cpp={metrics['worst_cpp']:.6e} " - f"n_over_atol={metrics['n_over_atol']}/{metrics['n_total']}") - - -# Define all decoder tensor cases with their tolerances -DECODER_TENSORS = [ - # Inputs (should be exact or near-exact) - ("ddec_query_embed", 1e-4), - ("ddec_ref_pts_raw", 1e-4), - ("ddec_presence_token", 1e-4), - ("ddec_ref_boxes_init", 1e-4), - - # Layer 0 intermediates - ("ddec_query_sine_0", 1e-4), - ("ddec_query_pos_0", 1e-4), - ("ddec_rpb_mask_0", 2e-4), - ("ddec_layer0_after_sa", 1e-4), - ("ddec_layer0_after_text_ca", 1e-4), - ("ddec_layer0_after_img_ca", 1e-4), - ("ddec_layer0_full_out", 1e-4), - ("ddec_layer0_presence", 1e-4), - - # Per-layer outputs (6 layers) - ("ddec_layer0_out", 1e-4), - ("ddec_layer0_refboxes", 1e-4), - ("ddec_layer1_out", 1e-4), - ("ddec_layer1_refboxes", 1e-4), - ("ddec_layer2_out", 2e-4), - ("ddec_layer2_refboxes", 1e-4), - ("ddec_layer3_out", 2e-4), - ("ddec_layer3_refboxes", 1e-4), - ("ddec_layer4_out", 2e-4), - ("ddec_layer4_refboxes", 1e-4), - ("ddec_layer5_out", 2e-4), - ("ddec_layer5_refboxes", 1e-4), - - # Final outputs - ("ddec_normed_output", 2e-4), - ("ddec_pred_boxes", 1e-4), - ("ddec_presence_logit", 1e-4), - - # Scoring - ("scoring_prompt_mlp_out", 1e-4), - ("scoring_pooled", 1e-4), - ("scoring_proj_pooled", 1e-4), - ("scoring_proj_hs", 1e-4), - ("scoring_class_scores", 1e-4), -] - - -def main(): - ref_dir = sys.argv[1] if len(sys.argv) > 1 else REF_DIR - cpp_dir = sys.argv[2] if len(sys.argv) > 2 else CPP_DIR - - print(f"Reference dir: {ref_dir}") - print(f"C++ output dir: {cpp_dir}") - print() - - n_pass = 0 - n_fail = 0 - n_skip = 0 - - # ═══ Section 1: Decoder Inputs ═══ - print("═══ Decoder Inputs (should be exact match) ═══") - for name, atol in DECODER_TENSORS[:4]: - metrics, status = compare(name, os.path.join(ref_dir, name), os.path.join(cpp_dir, name), atol) - print_result(name, metrics, status, atol) - if "PASS" in status: - n_pass += 1 - elif "FAIL" in status: - n_fail += 1 - else: - n_skip += 1 - - # ═══ Section 2: Decoder Layer 0 Intermediates ═══ - print("\n═══ Decoder Layer 0 Intermediates ═══") - for name, atol in DECODER_TENSORS[4:12]: - metrics, status = compare(name, os.path.join(ref_dir, name), os.path.join(cpp_dir, name), atol) - print_result(name, metrics, status, atol) - if "PASS" in status: - n_pass += 1 - elif "FAIL" in status: - n_fail += 1 - else: - n_skip += 1 - - # ═══ Section 3: Per-Layer Outputs ═══ - print("\n═══ Decoder Per-Layer Outputs (6 layers) ═══") - for name, atol in DECODER_TENSORS[12:24]: - metrics, status = compare(name, os.path.join(ref_dir, name), os.path.join(cpp_dir, name), atol) - print_result(name, metrics, status, atol) - if "PASS" in status: - n_pass += 1 - elif "FAIL" in status: - n_fail += 1 - else: - n_skip += 1 - - # ═══ Section 4: Final Decoder Outputs ═══ - print("\n═══ Final Decoder Outputs ═══") - for name, atol in DECODER_TENSORS[24:27]: - metrics, status = compare(name, os.path.join(ref_dir, name), os.path.join(cpp_dir, name), atol) - print_result(name, metrics, status, atol) - if "PASS" in status: - n_pass += 1 - elif "FAIL" in status: - n_fail += 1 - else: - n_skip += 1 - - # ═══ Section 5: Scoring ═══ - print("\n═══ DotProductScoring ═══") - for name, atol in DECODER_TENSORS[27:]: - metrics, status = compare(name, os.path.join(ref_dir, name), os.path.join(cpp_dir, name), atol) - print_result(name, metrics, status, atol) - if "PASS" in status: - n_pass += 1 - elif "FAIL" in status: - n_fail += 1 - else: - n_skip += 1 - - # ═══ Summary ═══ - print(f"\n{'='*70}") - total = n_pass + n_fail - print(f"SUMMARY: {n_pass}/{total} PASS, {n_fail}/{total} FAIL, {n_skip} SKIP") - if n_fail == 0 and n_skip == 0: - print("ALL DECODER TENSORS PASS ✓") - return 0 if n_fail == 0 else 1 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/tests/compare_f32_vs_f16.py b/tests/compare_f32_vs_f16.py deleted file mode 100644 index 684e932..0000000 --- a/tests/compare_f32_vs_f16.py +++ /dev/null @@ -1,105 +0,0 @@ -#!/usr/bin/env python3 -"""Side-by-side comparison: F32 vs F16 mask decoder precision against Python reference.""" -import numpy as np -import os - -REF_BASE = "tests/ref_phase6" -F32_BASE = "tests/ref_phase6/cpp_out_phase6" -F16_BASE = "tests/ref_phase6/cpp_out_phase6_f16/cpp_out_phase6" - -CASES = ["point_single", "box_only", "box_and_points"] - -TENSORS = [ - "sam_pe_sparse", "sam_pe_dense", "sam_pe_image_pe", - "sam_dec_image_feats", "sam_dec_tokens_initial", - "sam_dec_block0_queries", "sam_dec_block0_keys", - "sam_dec_block1_queries", "sam_dec_block1_keys", - "sam_dec_final_queries", - "sam_dec_feat_s1_proj", "sam_dec_feat_s0_proj", - "sam_dec_upscaled", "sam_dec_mask_tokens", - "sam_dec_masks", "sam_dec_iou", "sam_dec_obj_score", "sam_dec_sam_token", -] - -def load(path): - with open(path + ".shape") as f: - shape = [int(x) for x in f.read().strip().split(",") if x] - return np.fromfile(path + ".bin", dtype=np.float32), shape - -def metrics(ref, cpp): - diff = np.abs(ref - cpp) - eps = 1e-8 - return { - "mae": float(diff.mean()), - "max": float(diff.max()), - "rel": float((diff / (np.abs(ref) + eps)).mean()), - "cos": float(np.dot(ref, cpp) / (np.linalg.norm(ref) * np.linalg.norm(cpp) + 1e-12)), - "p95": float(np.percentile(diff, 95)), - "p99": float(np.percentile(diff, 99)), - } - -# Aggregate across all cases -all_f32 = {} -all_f16 = {} - -for case in CASES: - for t in TENSORS: - ref_path = os.path.join(REF_BASE, case, t) - f32_path = os.path.join(F32_BASE, case, t) - f16_path = os.path.join(F16_BASE, case, t) - if not all(os.path.exists(p + ".bin") for p in [ref_path, f32_path, f16_path]): - continue - ref_data, _ = load(ref_path) - f32_data, _ = load(f32_path) - f16_data, _ = load(f16_path) - if ref_data.size != f32_data.size or ref_data.size != f16_data.size: - continue - key = (case, t) - all_f32[key] = metrics(ref_data, f32_data) - all_f16[key] = metrics(ref_data, f16_data) - -# Print per-tensor summary (worst across cases) -print(f"{'Tensor':40s} | {'F32 max':>10s} {'F32 mae':>10s} {'F32 cos':>12s} | {'F16 max':>10s} {'F16 mae':>10s} {'F16 cos':>12s} | {'Ratio max':>10s}") -print("-" * 140) - -for t in TENSORS: - f32_worst_max = 0 - f32_worst_mae = 0 - f32_worst_cos = 1.0 - f16_worst_max = 0 - f16_worst_mae = 0 - f16_worst_cos = 1.0 - found = False - for case in CASES: - key = (case, t) - if key in all_f32: - found = True - f32_worst_max = max(f32_worst_max, all_f32[key]["max"]) - f32_worst_mae = max(f32_worst_mae, all_f32[key]["mae"]) - f32_worst_cos = min(f32_worst_cos, all_f32[key]["cos"]) - f16_worst_max = max(f16_worst_max, all_f16[key]["max"]) - f16_worst_mae = max(f16_worst_mae, all_f16[key]["mae"]) - f16_worst_cos = min(f16_worst_cos, all_f16[key]["cos"]) - if not found: - continue - ratio = f16_worst_max / f32_worst_max if f32_worst_max > 0 else float('inf') - print(f"{t:40s} | {f32_worst_max:10.2e} {f32_worst_mae:10.2e} {f32_worst_cos:12.8f} | {f16_worst_max:10.2e} {f16_worst_mae:10.2e} {f16_worst_cos:12.8f} | {ratio:10.1f}x") - -# Print per-case detail for key tensors -print(f"\n{'='*120}") -print("Detailed per-case metrics for key decoder tensors") -print(f"{'='*120}") -key_tensors = ["sam_dec_block0_queries", "sam_dec_block1_keys", "sam_dec_final_queries", - "sam_dec_upscaled", "sam_dec_masks", "sam_dec_iou", "sam_dec_obj_score"] - -for case in CASES: - print(f"\n Case: {case}") - print(f" {'Tensor':34s} | {'':5s} {'mae':>10s} {'max':>10s} {'rel':>10s} {'cos':>12s} {'p95':>10s} {'p99':>10s}") - print(f" {'-'*105}") - for t in key_tensors: - key = (case, t) - if key not in all_f32: - continue - f32 = all_f32[key] - f16 = all_f16[key] - print(f" {t:34s} | {'f32':5s} {f32['mae']:10.2e} {f32['max']:10.2e} {f32['rel']:10.2e} {f32['cos']:12.8f} {f32['p95']:10.2e} {f32['p99']:10.2e}") - print(f" {'':34s} | {'f16':5s} {f16['mae']:10.2e} {f16['max']:10.2e} {f16['rel']:10.2e} {f16['cos']:12.8f} {f16['p95']:10.2e} {f16['p99']:10.2e}") diff --git a/tests/compare_mask_decoder.py b/tests/compare_mask_decoder.py deleted file mode 100644 index d5f89ac..0000000 --- a/tests/compare_mask_decoder.py +++ /dev/null @@ -1,165 +0,0 @@ -#!/usr/bin/env python3 -""" -Comprehensive comparison of SAM3 mask decoder tensors: Python reference vs C++ output. -Reports mae, max, rel, cos, p95, p99 for every tensor in every test case. -""" -import os -import sys -import numpy as np - -import argparse -_parser = argparse.ArgumentParser() -_parser.add_argument("--cpp-base", default="tests/ref_phase6/cpp_out_phase6") -_parser.add_argument("--ref-base", default="tests/ref_phase6") -_args = _parser.parse_args() - -REF_BASE = _args.ref_base -CPP_BASE = _args.cpp_base - -CASES = ["point_single", "box_only", "box_and_points"] - -TENSORS = [ - "sam_pe_sparse", - "sam_pe_dense", - "sam_pe_image_pe", - "sam_dec_image_feats", - "sam_dec_tokens_initial", - "sam_dec_block0_queries", - "sam_dec_block0_keys", - "sam_dec_block1_queries", - "sam_dec_block1_keys", - "sam_dec_final_queries", - "sam_dec_feat_s1_proj", - "sam_dec_feat_s0_proj", - "sam_dec_upscaled", - "sam_dec_mask_tokens", - "sam_dec_masks", - "sam_dec_iou", - "sam_dec_obj_score", - "sam_dec_sam_token", -] - - -def load_tensor(path): - with open(path + ".shape") as f: - shape = [int(x) for x in f.read().strip().split(",") if x] - data = np.fromfile(path + ".bin", dtype=np.float32) - return data, shape - - -def compare(name, ref_path, cpp_path): - ref_data, ref_shape = load_tensor(ref_path) - cpp_data, cpp_shape = load_tensor(cpp_path) - - if ref_data.size != cpp_data.size: - return { - "status": "FAIL", - "note": f"SIZE MISMATCH ref={ref_data.size} cpp={cpp_data.size}", - "ref_shape": ref_shape, - "cpp_shape": cpp_shape, - } - - diff = np.abs(ref_data - cpp_data) - eps = 1e-8 - max_d = float(diff.max()) - mae = float(diff.mean()) - rel_err = diff / (np.abs(ref_data) + eps) - mean_rel = float(rel_err.mean()) - cos = float( - np.dot(ref_data, cpp_data) - / (np.linalg.norm(ref_data) * np.linalg.norm(cpp_data) + 1e-12) - ) - p95 = float(np.percentile(diff, 95)) - p99 = float(np.percentile(diff, 99)) - - worst_idx = int(np.argmax(diff)) - worst_ref = float(ref_data[worst_idx]) - worst_cpp = float(cpp_data[worst_idx]) - - return { - "status": "OK", - "ref_shape": ref_shape, - "cpp_shape": cpp_shape, - "mae": mae, - "max": max_d, - "rel": mean_rel, - "cos": cos, - "p95": p95, - "p99": p99, - "worst_idx": worst_idx, - "worst_ref": worst_ref, - "worst_cpp": worst_cpp, - "n_elements": ref_data.size, - } - - -def main(): - all_pass = True - - for case in CASES: - ref_dir = os.path.join(REF_BASE, case) - cpp_dir = os.path.join(CPP_BASE, case) - - print(f"\n{'='*80}") - print(f" Case: {case}") - print(f"{'='*80}") - - if not os.path.isdir(ref_dir): - print(f" [SKIP] Reference dir missing: {ref_dir}") - continue - if not os.path.isdir(cpp_dir): - print(f" [SKIP] C++ output dir missing: {cpp_dir}") - continue - - for tensor_name in TENSORS: - ref_path = os.path.join(ref_dir, tensor_name) - cpp_path = os.path.join(cpp_dir, tensor_name) - - if not os.path.exists(ref_path + ".bin"): - print(f" [SKIP] {tensor_name:40s} missing ref") - continue - if not os.path.exists(cpp_path + ".bin"): - print(f" [SKIP] {tensor_name:40s} missing cpp") - continue - - r = compare(tensor_name, ref_path, cpp_path) - - if r["status"] != "OK": - print(f" [FAIL] {tensor_name:40s} {r['note']}") - all_pass = False - continue - - # Determine pass/fail thresholds - is_fail = False - if r["cos"] < 0.999: - is_fail = True - if r["max"] > 0.1: - is_fail = True - - status = "FAIL" if is_fail else "PASS" - if is_fail: - all_pass = False - - print( - f" [{status}] {tensor_name:40s} shape={r['ref_shape']}" - ) - print( - f" mae={r['mae']:.6e} max={r['max']:.6e} " - f"rel={r['rel']:.6e} cos={r['cos']:.8f} " - f"p95={r['p95']:.6e} p99={r['p99']:.6e}" - ) - if is_fail: - print( - f" worst_idx={r['worst_idx']} " - f"ref={r['worst_ref']:.6e} cpp={r['worst_cpp']:.6e} " - f"n_elements={r['n_elements']}" - ) - - print(f"\n{'='*80}") - print(f" Overall: {'ALL PASS' if all_pass else 'FAILURES DETECTED'}") - print(f"{'='*80}") - return 0 if all_pass else 1 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/tests/compare_preprocessed_float.py b/tests/compare_preprocessed_float.py deleted file mode 100644 index 3268f51..0000000 --- a/tests/compare_preprocessed_float.py +++ /dev/null @@ -1,45 +0,0 @@ -#!/usr/bin/env python3 -"""Compare C++ preprocessed float CHW against Python preprocessed float CHW.""" -import numpy as np - -ref_dir = "tests/ref_phase3" - -# Load Python reference: [1, 3, 1008, 1008] NCHW -py = np.fromfile(f"{ref_dir}/preprocessed.bin", dtype=np.float32) -py_shape = [int(x) for x in open(f"{ref_dir}/preprocessed.shape").read().strip().split(",")] -py = py.reshape(py_shape) -print(f"Python preprocessed: shape={py.shape}") - -# Load C++ preprocessed: [1, 3, 1008, 1008] NCHW -cpp = np.fromfile(f"{ref_dir}/cpp_preprocessed.bin", dtype=np.float32) -cpp_shape = [int(x) for x in open(f"{ref_dir}/cpp_preprocessed.shape").read().strip().split(",")] -cpp = cpp.reshape(cpp_shape) -print(f"C++ preprocessed: shape={cpp.shape}") - -# Compare -diff = np.abs(py - cpp) -print(f"\n═══ Preprocessed Float Comparison ═══") -print(f" max_diff: {diff.max():.6e}") -print(f" mean_diff: {diff.mean():.6e}") -cos = np.dot(py.flatten(), cpp.flatten()) / (np.linalg.norm(py.flatten()) * np.linalg.norm(cpp.flatten())) -print(f" cosine: {cos:.10f}") -n_diff = (diff > 0).sum() -print(f" elements differing (>0): {n_diff} / {diff.size} ({100*n_diff/diff.size:.2f}%)") - -# Since normalization is (x/255 - 0.5)/0.5, a 1 uint8 diff = 1/255 * 2 = 0.00784 -expected_1px = 1.0 / 255.0 / 0.5 # 0.00784 -print(f"\n Expected diff for 1 uint8 pixel: {expected_1px:.6f}") -n_1px = (diff > expected_1px * 0.5).sum() -print(f" Elements with diff > 0.5 pixel: {n_1px} ({100*n_1px/diff.size:.2f}%)") - -# Look at where biggest differences are -worst = np.unravel_index(np.argmax(diff), diff.shape) -print(f"\n Worst at {worst}: py={py[worst]:.6f}, cpp={cpp[worst]:.6f}, diff={diff[worst]:.6f}") -print(f" That's {diff[worst] / expected_1px:.2f} uint8 pixel values") - -# Compare uint8 values (undo normalization) -py_uint8 = np.round((py * 0.5 + 0.5) * 255).astype(np.uint8) -cpp_uint8 = np.round((cpp * 0.5 + 0.5) * 255).astype(np.uint8) -uint8_diff = np.abs(py_uint8.astype(np.int16) - cpp_uint8.astype(np.int16)) -print(f"\n Recovered uint8 max pixel diff: {uint8_diff.max()}") -print(f" Recovered uint8 n_diff: {(uint8_diff > 0).sum()}") diff --git a/tests/compare_preprocessing.py b/tests/compare_preprocessing.py deleted file mode 100644 index b471280..0000000 --- a/tests/compare_preprocessing.py +++ /dev/null @@ -1,118 +0,0 @@ -#!/usr/bin/env python3 -"""Compare C++ vs Python image preprocessing. - -Replicates the C++ sam3_preprocess_image() in Python and compares -against the torchvision pipeline used by SAM3 Python. -""" -import numpy as np -import sys -from PIL import Image -import torch -from torchvision.transforms import v2 - - -def cpp_preprocess(image_path, img_size=1008): - """Replicate C++ sam3_preprocess_image exactly. - - C++ does: - 1. Bilinear resize to img_size x img_size (half-pixel center mapping) - 2. Quantize back to uint8 with rounding - 3. Normalize: (pixel/255.0 - 0.5) / 0.5 - 4. Layout: CHW - """ - img = Image.open(image_path).convert("RGB") - src = np.array(img, dtype=np.float32) # [H, W, 3] - src_h, src_w = src.shape[:2] - - # C++ bilinear resize (half-pixel center mapping) - dst = np.zeros((img_size, img_size, 3), dtype=np.uint8) - sx = src_w / img_size - sy = src_h / img_size - - for y in range(img_size): - fy = (y + 0.5) * sy - 0.5 - y0 = max(0, int(fy)) - y1 = min(src_h - 1, y0 + 1) - wy = fy - y0 - for x in range(img_size): - fx = (x + 0.5) * sx - 0.5 - x0 = max(0, int(fx)) - x1 = min(src_w - 1, x0 + 1) - wx = fx - x0 - for c in range(3): - v = ((1 - wy) * ((1 - wx) * src[y0, x0, c] + wx * src[y0, x1, c]) + - wy * ((1 - wx) * src[y1, x0, c] + wx * src[y1, x1, c])) - dst[y, x, c] = min(255, max(0, int(v + 0.5))) - - # Normalize and convert to CHW - result = np.zeros((3, img_size, img_size), dtype=np.float32) - for c in range(3): - for y in range(img_size): - for x in range(img_size): - v = dst[y, x, c] / 255.0 - result[c, y, x] = (v - 0.5) / 0.5 - - return result, dst - - -def python_preprocess(image_path, img_size=1008): - """SAM3 Python preprocessing pipeline.""" - img = Image.open(image_path).convert("RGB") - transform = v2.Compose([ - v2.ToDtype(torch.uint8, scale=True), - v2.Resize(size=(img_size, img_size)), - v2.ToDtype(torch.float32, scale=True), - v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), - ]) - img_tensor = v2.functional.to_image(img) - return transform(img_tensor).numpy() # [3, H, W] - - -if __name__ == "__main__": - image_path = sys.argv[1] if len(sys.argv) > 1 else "tests/test_random.jpg" - - print(f"Comparing preprocessing for: {image_path}") - print() - - # Python preprocessing (torchvision) - py = python_preprocess(image_path) - print(f"Python shape: {py.shape}") - print(f" min={py.min():.6f} max={py.max():.6f} mean={py.mean():.6f}") - - # C++ preprocessing (replicated) - print("\nReplicating C++ preprocessing (this is slow in Python)...") - cpp, cpp_uint8 = cpp_preprocess(image_path) - print(f"C++ shape: {cpp.shape}") - print(f" min={cpp.min():.6f} max={cpp.max():.6f} mean={cpp.mean():.6f}") - - # Compare - diff = np.abs(py - cpp) - print(f"\n═══ Comparison ═══") - print(f" max_diff: {diff.max():.6e}") - print(f" mean_diff: {diff.mean():.6e}") - cos = np.dot(py.flatten(), cpp.flatten()) / (np.linalg.norm(py.flatten()) * np.linalg.norm(cpp.flatten())) - print(f" cosine: {cos:.10f}") - - # Pixel-level differences - n_diff = (diff > 1e-6).sum() - print(f" pixels differing: {n_diff} / {diff.size} ({100*n_diff/diff.size:.2f}%)") - - n_diff_1pct = (diff > 0.01).sum() - print(f" pixels with |err| > 0.01: {n_diff_1pct}") - - # Sample some differences - if diff.max() > 1e-6: - worst = np.unravel_index(np.argmax(diff), diff.shape) - print(f"\n Worst pixel at {worst}: py={py[worst]:.6f}, cpp={cpp[worst]:.6f}") - - # Check if the difference is due to resize or normalization - # PIL resize - img = Image.open(image_path).convert("RGB") - pil_resized = img.resize((1008, 1008), Image.BILINEAR) - pil_arr = np.array(pil_resized, dtype=np.uint8) - - # Compare uint8 resize outputs - uint8_diff = np.abs(pil_arr.astype(np.int16) - cpp_uint8.astype(np.int16)) - print(f"\n Resize comparison (uint8):") - print(f" max pixel diff: {uint8_diff.max()}") - print(f" pixels differing: {(uint8_diff > 0).sum()} / {uint8_diff.size}") diff --git a/tests/compare_resize.py b/tests/compare_resize.py deleted file mode 100644 index 24afd52..0000000 --- a/tests/compare_resize.py +++ /dev/null @@ -1,100 +0,0 @@ -#!/usr/bin/env python3 -"""Compare C++ bilinear resize against Python torchvision resize at uint8 level. - -Replicates the exact C++ sam3_resize_bilinear() in Python, then compares -against torchvision v2.Resize output. -""" -import numpy as np -from PIL import Image -import torch -from torchvision.transforms import v2 - -image_path = "tests/test_random.jpg" -img_size = 1008 - -# Load original image -img = Image.open(image_path).convert("RGB") -src = np.array(img, dtype=np.float32) # [H, W, 3] -src_h, src_w = src.shape[:2] -print(f"Original: {src_w}x{src_h}") - -# ═══ Method 1: C++ bilinear resize (replicated in Python) ═══ -print("\nReplicating C++ resize...") -dst_cpp = np.zeros((img_size, img_size, 3), dtype=np.uint8) -sx = src_w / img_size -sy = src_h / img_size - -for y in range(img_size): - fy = (y + 0.5) * sy - 0.5 - y0 = max(0, int(fy)) - y1 = min(src_h - 1, y0 + 1) - wy = np.float32(fy - y0) # Use f32 to match C++ - for x in range(img_size): - fx = (x + 0.5) * sx - 0.5 - x0 = max(0, int(fx)) - x1 = min(src_w - 1, x0 + 1) - wx = np.float32(fx - x0) - for c in range(3): - v = np.float32((1 - wy) * ((1 - wx) * src[y0, x0, c] + wx * src[y0, x1, c]) + - wy * ((1 - wx) * src[y1, x0, c] + wx * src[y1, x1, c])) - dst_cpp[y, x, c] = min(255, max(0, int(v + 0.5))) - -# ═══ Method 2: Python torchvision resize ═══ -img_tensor = v2.functional.to_image(img) # uint8 CHW tensor -step1 = v2.ToDtype(torch.uint8, scale=True)(img_tensor) -step2 = v2.Resize(size=(img_size, img_size))(step1) -dst_py = step2.permute(1, 2, 0).numpy() # [H, W, 3] uint8 - -# ═══ Compare ═══ -print("\n═══ C++ resize vs Python torchvision resize (uint8) ═══") -diff = np.abs(dst_cpp.astype(np.int16) - dst_py.astype(np.int16)) -print(f" max pixel diff: {diff.max()}") -print(f" mean pixel diff: {diff.mean():.6f}") -n_diff = (diff > 0).sum() -print(f" pixels differing: {n_diff} / {diff.size} ({100*n_diff/diff.size:.2f}%)") - -# Histogram of differences -for threshold in [0, 1, 2, 3]: - n = (diff > threshold).sum() - print(f" pixels with diff > {threshold}: {n}") - -# ═══ Check: does the float normalization match? ═══ -print("\n═══ Float normalization check ═══") - -# C++ normalization: (pixel / 255.0 - 0.5) / 0.5 -cpp_float = (dst_cpp.astype(np.float32) / 255.0 - 0.5) / 0.5 # [H, W, 3] -cpp_chw = cpp_float.transpose(2, 0, 1) # [3, H, W] - -# Python: v2.ToDtype(float32, scale=True) then Normalize(0.5, 0.5) -step3 = v2.ToDtype(torch.float32, scale=True)(step2) # scale=True divides by 255 -step4 = v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])(step3) -py_chw = step4.numpy() - -# If the uint8 images match, the float should also match -float_diff = np.abs(cpp_chw - py_chw) -print(f" max float diff: {float_diff.max():.6e}") -print(f" mean float diff: {float_diff.mean():.6e}") - -# ═══ Check: what if we use the SAME uint8 but different normalization? ═══ -# Python's ToDtype(float32, scale=True) does pixel / 255.0 -# C++ does pixel / 255.0 -# Then Python's Normalize does (x - 0.5) / 0.5 -# C++ does (v - 0.5) / 0.5 -# These should be identical -print("\n═══ Normalization formula check ═══") -# Take a single pixel value and check -test_val = np.uint8(128) -cpp_result = np.float32((np.float32(test_val) / np.float32(255.0) - np.float32(0.5)) / np.float32(0.5)) -py_result = np.float32((np.float32(test_val) / np.float32(255.0) - np.float32(0.5)) / np.float32(0.5)) -print(f" test pixel=128: cpp={cpp_result:.10f}, py={py_result:.10f}, diff={abs(cpp_result-py_result):.10e}") - -# Conclusion -print("\n═══ Conclusion ═══") -if diff.max() <= 1: - print(" C++ and Python resizes differ by at most 1 uint8 value.") - print(" This is a bilinear rounding difference, not a bug.") - print(" To make them identical, the C++ resize must replicate") - print(" the exact torchvision/PIL rounding behavior.") -else: - print(f" WARNING: C++ and Python resizes differ by up to {diff.max()} uint8 values!") - print(" This suggests an algorithmic difference, not just rounding.") diff --git a/tests/compare_sam2_video.py b/tests/compare_sam2_video.py deleted file mode 100644 index 89612a8..0000000 --- a/tests/compare_sam2_video.py +++ /dev/null @@ -1,122 +0,0 @@ -#!/usr/bin/env python3 -""" -Compare C++ vs Python SAM2 video tracking outputs. - -Compares: -1. Backbone features (neck_trk_2) — are images encoded the same? -2. Output masks — do they segment the same region? -3. Memory-conditioned features — is memory attention working? -""" -import numpy as np -import os, sys -from PIL import Image - -py_dir = "/tmp/debug_sam2_pipeline" -cpp_dir = "/tmp/debug_sam2_cpp" - -def load_tensor(path): - data = np.fromfile(path + ".bin", dtype=np.float32) - if os.path.exists(path + ".shape"): - with open(path + ".shape") as f: - shape = [int(x) for x in f.read().strip().split(",") if x] - return data.reshape(shape) if shape else data - return data - -def compare(name, ref, cpp, atol=1e-3): - ref_f = ref.flatten() - cpp_f = cpp.flatten() - if ref_f.size != cpp_f.size: - print(f" FAIL {name:50s} SIZE MISMATCH ref={ref_f.size} cpp={cpp_f.size}") - return False - diff = np.abs(ref_f - cpp_f) - max_d = diff.max() - mean_d = diff.mean() - cos = np.dot(ref_f, cpp_f) / (np.linalg.norm(ref_f) * np.linalg.norm(cpp_f) + 1e-12) - p95 = np.percentile(diff, 95) - p99 = np.percentile(diff, 99) - status = "PASS" if cos > 0.99 else "FAIL" - print(f" {status} {name:50s}") - print(f" mae={mean_d:.6e} max={max_d:.6e} cos={cos:.8f} p95={p95:.6e} p99={p99:.6e}") - return status == "PASS" - -def mask_iou(m1, m2): - """Compute IoU between two binary masks.""" - intersection = np.logical_and(m1, m2).sum() - union = np.logical_or(m1, m2).sum() - return intersection / max(union, 1) - -print("=" * 80) -print("SAM2 Video Tracking: Python vs C++ Comparison") -print("=" * 80) - -# Compare vision features (backbone output) -# Python vision_feat_0 is [4096, 1, 256] (HW, B, C) = flattened (64, 64, 256) -# C++ neck_trk_2 is [256, 64, 64] in ggml -print("\n--- Backbone Features (neck_trk_2 / vision_feat_0) ---") -for fi in range(6): - py_path = os.path.join(py_dir, f"f{fi}_vision_feat_0") - cpp_path = os.path.join(cpp_dir, f"f{fi}_neck_trk_2") - if not os.path.exists(py_path + ".bin") or not os.path.exists(cpp_path + ".bin"): - print(f" Frame {fi}: SKIP (missing tensors)") - continue - - py_data = load_tensor(py_path) # [4096, 1, 256] - cpp_data = load_tensor(cpp_path) # [256, 64, 64] in ggml - - print(f" Frame {fi}: py={py_data.shape} ({py_data.size} floats), cpp={cpp_data.shape} ({cpp_data.size} floats)") - - # The Python tensor is [HW, B, C] = [4096, 1, 256] - # The C++ tensor is ggml [D, H, W] = [256, 64, 64] - # Both have 256*64*64 = 1048576 floats - - # Flatten and compare (cosine similarity should be near 1 if they match) - compare(f"frame{fi}_backbone", py_data, cpp_data) - -# Compare output masks -print("\n--- Output Masks ---") -for fi in range(1, 6): - py_mask_path = os.path.join(py_dir, f"f{fi}_output_mask") - cpp_mask_path = os.path.join(cpp_dir, f"f{fi}_output_mask_binary") - - if not os.path.exists(py_mask_path + ".bin") or not os.path.exists(cpp_mask_path + ".bin"): - print(f" Frame {fi}: SKIP") - continue - - py_mask = load_tensor(py_mask_path) # [1, 540, 960] logits - cpp_mask = load_tensor(cpp_mask_path) # [540, 960] binary - - py_binary = (py_mask.flatten() > 0.0).astype(float) - cpp_binary = cpp_mask.flatten() - - iou = mask_iou(py_binary > 0.5, cpp_binary > 0.5) - py_fg = (py_binary > 0.5).sum() - cpp_fg = (cpp_binary > 0.5).sum() - - status = "PASS" if iou > 0.5 else "FAIL" - print(f" {status} frame{fi}_mask IoU={iou:.4f} py_fg={int(py_fg)} cpp_fg={int(cpp_fg)}") - -# Compare memory-conditioned features -print("\n--- Memory-Conditioned Features ---") -for fi in range(6): - py_path = os.path.join(py_dir, f"f{fi}_mem_conditioned") - if not os.path.exists(py_path + ".bin"): - continue - py_data = load_tensor(py_path) # [1, 256, 64, 64] - print(f" Frame {fi}: py_mem_conditioned shape={py_data.shape} " - f"range=[{py_data.min():.4f}, {py_data.max():.4f}]") - -# Compare SAM decoder outputs -print("\n--- SAM Decoder Outputs ---") -for fi in range(6): - for name in ["sam_pred_masks", "sam_ious", "sam_obj_ptr", "sam_obj_score"]: - path = os.path.join(py_dir, f"f{fi}_{name}") - if os.path.exists(path + ".bin"): - data = load_tensor(path) - print(f" Frame {fi} {name:20s}: shape={data.shape} " - f"range=[{data.min():.4f}, {data.max():.4f}]") - -# Print summary -print("\n--- Summary ---") -print("Python ref: 2000-2200 fg pixels/frame, IoU scores 0.88-0.90") -print("If C++ masks have similar fg counts but different locations, memory attention is wrong.") -print("If C++ masks are reasonable, the issue may be in longer sequences (memory drift).") diff --git a/tests/compare_tensors.py b/tests/compare_tensors.py deleted file mode 100644 index c3fc608..0000000 --- a/tests/compare_tensors.py +++ /dev/null @@ -1,57 +0,0 @@ -#!/usr/bin/env python3 -"""Compare tensors from Python reference vs C++ output.""" -import numpy as np -import sys, os - -def load_tensor(path): - data = np.fromfile(path + ".bin", dtype=np.float32) - shape_file = path + ".shape" - if os.path.exists(shape_file): - with open(shape_file) as f: - shape = [int(x) for x in f.read().strip().split(",") if x] - return data.reshape(shape) if shape else data - return data - -def compare(name, ref_path, cpp_path, atol=1e-4): - ref = load_tensor(ref_path) if isinstance(ref_path, str) else ref_path - cpp = load_tensor(cpp_path) if isinstance(cpp_path, str) else cpp_path - ref_f = ref.flatten() - cpp_f = cpp.flatten() - if ref_f.size != cpp_f.size: - print(f" FAIL {name:40s} SIZE MISMATCH ref={ref_f.size} cpp={cpp_f.size}") - return False - diff = np.abs(ref_f - cpp_f) - eps = 1e-8 - max_d = diff.max() - mean_d = diff.mean() - rel_err = (diff / (np.abs(ref_f) + eps)).mean() - cos = np.dot(ref_f, cpp_f) / (np.linalg.norm(ref_f) * np.linalg.norm(cpp_f) + 1e-12) - p95 = np.percentile(diff, 95) - p99 = np.percentile(diff, 99) - status = "PASS" if max_d < atol else "FAIL" - print(f" {status} {name:40s} shape={ref.shape}") - print(f" mae={mean_d:.6e} max={max_d:.6e} rel={rel_err:.6e} cos={cos:.8f} p95={p95:.6e} p99={p99:.6e}") - if status == "FAIL": - worst = np.argmax(diff) - print(f" worst idx={worst} ref={ref_f[worst]:.6e} cpp={cpp_f[worst]:.6e}") - return status == "PASS" - -if __name__ == "__main__": - ref_dir = sys.argv[1] # Python dump dir - cpp_dir = sys.argv[2] # C++ dump dir - atol = float(sys.argv[3]) if len(sys.argv) > 3 else 1e-4 - ref_tensors = {f[:-4] for f in os.listdir(ref_dir) if f.endswith(".bin")} - cpp_tensors = {f[:-4] for f in os.listdir(cpp_dir) if f.endswith(".bin")} - common = sorted(ref_tensors & cpp_tensors) - if not common: - print("No common tensors found!") - ref_only = sorted(ref_tensors - cpp_tensors) - cpp_only = sorted(cpp_tensors - ref_tensors) - if ref_only: print(f" Ref only: {ref_only[:10]}") - if cpp_only: print(f" Cpp only: {cpp_only[:10]}") - sys.exit(1) - n_pass = 0 - for name in common: - if compare(name, os.path.join(ref_dir, name), os.path.join(cpp_dir, name), atol): - n_pass += 1 - print(f"\n{n_pass}/{len(common)} passed (atol={atol})") diff --git a/tests/compare_vit_output.py b/tests/compare_vit_output.py deleted file mode 100644 index d08bef7..0000000 --- a/tests/compare_vit_output.py +++ /dev/null @@ -1,96 +0,0 @@ -#!/usr/bin/env python3 -"""Compare C++ and Python ViT outputs in detail to find outliers. - -The C++ dump stores data in ggml column-major order with shape ne[0], ne[1], ne[2]. -For a tensor with ne=[E, W, H], element (e, w, h) is at flat index e + w*E + h*E*W. -This is the REVERSE of numpy's row-major convention. -""" -import numpy as np -import sys -import os - -def load_ggml_tensor(path): - """Load a tensor dumped by sam3_dump_state_tensor. - Returns data in ggml flat order and the ggml shape (ne[0], ne[1], ne[2], ...). - """ - with open(path + ".shape") as f: - shape = [int(x) for x in f.read().strip().split(",")] - with open(path + ".bin", "rb") as f: - data = np.frombuffer(f.read(), dtype=np.float32).copy() - return data, shape - -def ggml_to_pytorch_bchw(data, ggml_shape): - """Convert ggml [E, W, H] flat data to PyTorch [1, C, H, W] array. - ggml element (e, w, h) at flat index e + w*E + h*E*W - maps to PyTorch [0, e, h, w]. - """ - E, W, H = ggml_shape[:3] - result = np.zeros((1, E, H, W), dtype=np.float32) - for h in range(H): - for w in range(W): - for e in range(E): - result[0, e, h, w] = data[e + w*E + h*E*W] - return result - -def load_pytorch_tensor(path): - """Load a tensor saved by the Python dumper (standard row-major).""" - with open(path + ".shape") as f: - shape = [int(x) for x in f.read().strip().split(",")] - with open(path + ".bin", "rb") as f: - data = np.frombuffer(f.read(), dtype=np.float32).copy() - return data.reshape(shape) - -ref_dir = sys.argv[1] if len(sys.argv) > 1 else "tests/ref_phase3" -cpp_dir = sys.argv[2] if len(sys.argv) > 2 else "tests/ref_phase3/cpp_out_phase3" - -# Load Python reference (BCHW) -py_vit = load_pytorch_tensor(f"{ref_dir}/vit_output_bchw") # [1, 1024, 72, 72] -print(f"Python ViT output: shape={py_vit.shape}, range=[{py_vit.min():.4f}, {py_vit.max():.4f}]") - -# Load C++ output (ggml: ne=[E=1024, W=72, H=72]) -cpp_data, cpp_shape = load_ggml_tensor(f"{cpp_dir}/vit_output") -print(f"C++ ViT output: ggml_shape={cpp_shape}, range=[{cpp_data.min():.4f}, {cpp_data.max():.4f}]") - -# Transpose properly -E, W, H = cpp_shape[:3] -print(f" Transposing ggml [E={E}, W={W}, H={H}] to PyTorch [1, {E}, {H}, {W}]...") -cpp_transposed = ggml_to_pytorch_bchw(cpp_data, cpp_shape) - -diff = np.abs(cpp_transposed - py_vit) -print(f"\nDifference stats:") -print(f" max_diff: {diff.max():.6f}") -print(f" mean_diff: {diff.mean():.6f}") -print(f" median: {np.median(diff):.6f}") -print(f" 99.9th %: {np.percentile(diff, 99.9):.6f}") -print(f" 99.99th %: {np.percentile(diff, 99.99):.6f}") - -# Find the worst elements -flat_diff = diff.flatten() -worst_indices = np.argsort(flat_diff)[-20:][::-1] -print(f"\nTop 20 worst mismatches:") -for idx in worst_indices: - b, c, h, w = np.unravel_index(idx, diff.shape) - py_val = py_vit[b, c, h, w] - cpp_val = cpp_transposed[b, c, h, w] - d = flat_diff[idx] - rel = d / max(abs(py_val), 1e-8) - print(f" [{b},{c:4d},{h:2d},{w:2d}] py={py_val:10.4f} cpp={cpp_val:10.4f} diff={d:.4f} rel={rel:.4f}") - -# Per-channel analysis -chan_max = diff[0].max(axis=(1, 2)) # [1024] -print(f"\nPer-channel max diff distribution:") -for threshold in [0.1, 0.5, 1.0, 2.0, 5.0, 10.0]: - n = (chan_max > threshold).sum() - print(f" channels with max_diff > {threshold:5.1f}: {n}/{E}") - -# Overall element analysis -for threshold in [0.01, 0.1, 0.5, 1.0, 5.0]: - n = (diff > threshold).sum() - print(f"Elements with diff > {threshold}: {n}/{diff.size} ({100*n/diff.size:.2f}%)") - -# Cosine similarity -dot = np.sum(cpp_transposed * py_vit) -norm_a = np.sqrt(np.sum(cpp_transposed**2)) -norm_b = np.sqrt(np.sum(py_vit**2)) -cosine = dot / (norm_a * norm_b) -print(f"\nCosine similarity: {cosine:.8f}") diff --git a/tests/dump_cpp_preprocessed.py b/tests/dump_cpp_preprocessed.py deleted file mode 100644 index 7f20663..0000000 --- a/tests/dump_cpp_preprocessed.py +++ /dev/null @@ -1,62 +0,0 @@ -#!/usr/bin/env python3 -"""Compare C++ vs Python image preprocessing. - -The C++ preprocessing: -1. Resize to 1008x1008 using bilinear interpolation (custom implementation) -2. Normalize: (pixel / 255.0 - 0.5) / 0.5 = pixel/127.5 - 1.0 - -The Python preprocessing (torchvision): -1. v2.Resize(size=(1008, 1008)) — uses PIL's bilinear -2. v2.Normalize(mean=[0.5,0.5,0.5], std=[0.5,0.5,0.5]) - -Differences come from the resize algorithm (C++ custom vs PIL bilinear). -This script replicates the C++ preprocessing in Python for verification. -""" -import numpy as np -import sys -from PIL import Image - -def cpp_style_preprocess(image_path, img_size=1008): - """Replicate C++ sam3_preprocess_image exactly.""" - img = Image.open(image_path).convert("RGB") - # PIL bilinear resize (this may differ from C++ implementation) - img_resized = img.resize((img_size, img_size), Image.BILINEAR) - pixels = np.array(img_resized, dtype=np.float32) # [H, W, 3] - - # Normalize: (pixel / 255.0 - 0.5) / 0.5 - pixels = (pixels / 255.0 - 0.5) / 0.5 - - # Convert HWC → CHW - chw = pixels.transpose(2, 0, 1) # [3, H, W] - return chw - -def load_tensor(path): - with open(path + ".shape") as f: - shape = [int(x) for x in f.read().strip().split(",") if x] - data = np.fromfile(path + ".bin", dtype=np.float32) - return data.reshape(shape) - -if __name__ == "__main__": - ref_dir = sys.argv[1] if len(sys.argv) > 1 else "tests/ref_phase3" - - # Load Python reference - ref = load_tensor(f"{ref_dir}/preprocessed") - print(f"Python ref shape: {ref.shape}") - print(f" min={ref.min():.6f} max={ref.max():.6f} mean={ref.mean():.6f}") - - # Also load raw float (before normalization) for debugging - try: - raw = load_tensor(f"{ref_dir}/preprocessed_raw_float") - print(f"Python raw float shape: {raw.shape}") - print(f" min={raw.min():.6f} max={raw.max():.6f} mean={raw.mean():.6f}") - except: - pass - - # Quick sanity: check normalization - # PyTorch Normalize: (x - 0.5) / 0.5 where x is in [0,1] - # So output range should be [-1, 1] - print(f"\n Expected range: [-1, 1]") - print(f" Actual range: [{ref.min():.6f}, {ref.max():.6f}]") - - # Check first few pixels for debugging - print(f"\n First 10 pixels (channel 0): {ref.flatten()[:10]}") diff --git a/tests/dump_ddec_from_package.py b/tests/dump_ddec_from_package.py deleted file mode 100644 index f230e1a..0000000 --- a/tests/dump_ddec_from_package.py +++ /dev/null @@ -1,411 +0,0 @@ -#!/usr/bin/env python3 -""" -Dump DETR decoder reference tensors directly from the SAM3 Python package. -Uses the actual TransformerDecoder.forward() from the package (not reimplemented). - -This script: -1. Builds the model using the SAM3 package -2. Loads the checkpoint -3. Runs just the decoder forward pass with known inputs -4. Dumps per-layer outputs for comparison with C++ - -Usage: - cd ~/Documents/sam3 - python3 /path/to/dump_ddec_from_package.py --checkpoint raw_weights/sam3.pt --outdir /tmp/ddec_ref -""" -import argparse -import math -import os -import sys -import importlib - -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F - - -def save_tensor(path, t): - """Save tensor as .bin + .shape files.""" - t = t.detach().cpu().float().contiguous() - os.makedirs(os.path.dirname(path) or ".", exist_ok=True) - with open(path + ".bin", "wb") as f: - f.write(t.numpy().tobytes()) - with open(path + ".shape", "w") as f: - f.write(",".join(str(d) for d in t.shape)) - print(f" saved {path} shape={list(t.shape)}") - - -def inverse_sigmoid(x, eps=1e-3): - x = x.clamp(min=0, max=1) - x1 = x.clamp(min=eps) - x2 = (1 - x).clamp(min=eps) - return torch.log(x1 / x2) - - -def gen_sineembed_for_position(pos_tensor, num_feats=256): - assert num_feats % 2 == 0 - num_feats = num_feats // 2 - scale = 2 * math.pi - dim_t = torch.arange(num_feats, dtype=torch.float32, device=pos_tensor.device) - dim_t = 10000 ** (2 * (torch.div(dim_t, 2, rounding_mode="floor")) / num_feats) - x_embed = pos_tensor[:, :, 0] * scale - y_embed = pos_tensor[:, :, 1] * scale - pos_x = x_embed[:, :, None] / dim_t - pos_y = y_embed[:, :, None] / dim_t - pos_x = torch.stack((pos_x[:, :, 0::2].sin(), pos_x[:, :, 1::2].cos()), dim=3).flatten(2) - pos_y = torch.stack((pos_y[:, :, 0::2].sin(), pos_y[:, :, 1::2].cos()), dim=3).flatten(2) - if pos_tensor.size(-1) == 2: - pos = torch.cat((pos_y, pos_x), dim=2) - elif pos_tensor.size(-1) == 4: - w_embed = pos_tensor[:, :, 2] * scale - pos_w = w_embed[:, :, None] / dim_t - pos_w = torch.stack((pos_w[:, :, 0::2].sin(), pos_w[:, :, 1::2].cos()), dim=3).flatten(2) - h_embed = pos_tensor[:, :, 3] * scale - pos_h = h_embed[:, :, None] / dim_t - pos_h = torch.stack((pos_h[:, :, 0::2].sin(), pos_h[:, :, 1::2].cos()), dim=3).flatten(2) - pos = torch.cat((pos_y, pos_x, pos_w, pos_h), dim=2) - else: - raise ValueError(f"Unknown pos_tensor shape(-1):{pos_tensor.size(-1)}") - return pos - - -def box_cxcywh_to_xyxy(x): - x_c, y_c, w, h = x.unbind(-1) - b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c + 0.5 * w), (y_c + 0.5 * h)] - return torch.stack(b, dim=-1) - - -def mlp_forward(x, sd, prefix, num_layers, activation="relu"): - for i in range(num_layers): - w = sd[f"{prefix}.layers.{i}.weight"] - b = sd[f"{prefix}.layers.{i}.bias"] - x = F.linear(x, w, b) - if i < num_layers - 1: - x = F.relu(x) - return x - - -def multihead_attention_forward(query, key, value, w_q, b_q, w_k, b_k, w_v, b_v, - w_out, b_out, num_heads, attn_mask=None, - key_padding_mask=None): - D = query.shape[-1] - HD = D // num_heads - Q = F.linear(query, w_q, b_q) - K = F.linear(key, w_k, b_k) - V = F.linear(value, w_v, b_v) - BS = Q.shape[0] - N_q = Q.shape[1] - N_kv = K.shape[1] - Q = Q.reshape(BS, N_q, num_heads, HD).permute(0, 2, 1, 3) - K = K.reshape(BS, N_kv, num_heads, HD).permute(0, 2, 1, 3) - V = V.reshape(BS, N_kv, num_heads, HD).permute(0, 2, 1, 3) - - combined_mask = attn_mask - if key_padding_mask is not None: - pad_bias = torch.zeros(BS, 1, 1, N_kv, dtype=Q.dtype, device=Q.device) - pad_bias = pad_bias.masked_fill(key_padding_mask[:, None, None, :], float("-inf")) - combined_mask = pad_bias if combined_mask is None else combined_mask + pad_bias - attn_output = F.scaled_dot_product_attention(Q, K, V, attn_mask=combined_mask) - attn_output = attn_output.permute(0, 2, 1, 3).reshape(BS, N_q, D) - out = F.linear(attn_output, w_out, b_out) - return out - - -def fused_mha_forward(query, key, value, in_proj_w, in_proj_b, out_w, out_b, - num_heads, attn_mask=None, key_padding_mask=None): - D = query.shape[-1] - w_q, w_k, w_v = in_proj_w[:D], in_proj_w[D:2*D], in_proj_w[2*D:] - b_q, b_k, b_v = in_proj_b[:D], in_proj_b[D:2*D], in_proj_b[2*D:] - return multihead_attention_forward(query, key, value, w_q, b_q, w_k, b_k, w_v, b_v, - out_w, out_b, num_heads, attn_mask, key_padding_mask) - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--checkpoint", required=True) - parser.add_argument("--outdir", default="/tmp/ddec_ref") - parser.add_argument("--fenc-input", default=None, - help="Directory containing fenc output to use as decoder input") - args = parser.parse_args() - - os.makedirs(args.outdir, exist_ok=True) - device = "cpu" - D = 256 - NH = 8 - NQ = 200 - H = 72 - - print("Loading checkpoint...") - ckpt = torch.load(args.checkpoint, map_location=device, weights_only=False) - if "model" in ckpt: - ckpt = ckpt["model"] - - ddec_prefix = "detector.transformer.decoder." - ddec_sd = {k[len("detector.transformer."):]: v for k, v in ckpt.items() - if k.startswith("detector.transformer.decoder.")} - - # Load fenc output as decoder input - fenc_dir = args.fenc_input - if fenc_dir is None: - # Use existing ref_phase5 data - script_dir = os.path.dirname(os.path.abspath(__file__)) - fenc_dir = os.path.join(script_dir, "ref_phase5") - - def load_tensor(path, dtype=np.float32): - shape = list(map(int, open(path + ".shape").read().strip().split(","))) - data = np.fromfile(path + ".bin", dtype=dtype).reshape(shape) - return torch.tensor(data) - - # Load fenc output and image PE - fenc_output_bf = load_tensor(os.path.join(fenc_dir, "fenc_output_bf")) - if fenc_output_bf.ndim == 3: - pass # [1, HW, D] batch-first - else: - # Try loading seq-first [HW, 1, D] and converting - fenc_output_sf = load_tensor(os.path.join(fenc_dir, "fenc_output")) - fenc_output_bf = fenc_output_sf.permute(1, 0, 2) # [1, HW, D] - - # Image PE - load and flatten - img_pe = load_tensor(os.path.join(fenc_dir, "img_pe_72")) # [1, 256, 72, 72] - img_pe_flat = img_pe.flatten(2).permute(0, 2, 1) # [1, 5184, 256] - - # Text features - text_features = load_tensor(os.path.join(fenc_dir, "text_features")) - if text_features.ndim == 3 and text_features.shape[0] == 32: - # [T, 1, D] seq-first - prompt_sf = text_features - prompt_bf = text_features.permute(1, 0, 2) - elif text_features.ndim == 2: - # [D, T] ggml layout -> convert - prompt_sf = text_features.T.unsqueeze(1) # [T, 1, D] - prompt_bf = prompt_sf.permute(1, 0, 2) - else: - prompt_sf = text_features - prompt_bf = text_features.permute(1, 0, 2) - - T = prompt_sf.shape[0] - - # Token IDs for text mask (saved as int32) - token_ids_path = os.path.join(fenc_dir, "token_ids") - if os.path.exists(token_ids_path + ".bin"): - token_ids = load_tensor(token_ids_path, dtype=np.int32).long().squeeze().tolist() - else: - # Assume all valid - token_ids = [1] * T - - txt_mask = torch.tensor([[tid == 0 for tid in token_ids]], dtype=torch.bool) - - print(f"fenc output: {list(fenc_output_bf.shape)}") - print(f"prompt: {list(prompt_sf.shape)}") - print(f"img PE: {list(img_pe_flat.shape)}") - - # ═══════════════════════════════════════════════════════════════════ - # DETR Decoder forward (matching Python package exactly) - # ═══════════════════════════════════════════════════════════════════ - print("\n=== DETR Decoder (from checkpoint weights) ===") - - with torch.no_grad(): - query_embed = ddec_sd["decoder.query_embed.weight"] # [NQ, D] - ref_pts_raw = ddec_sd["decoder.reference_points.weight"] # [NQ, 4] - reference_boxes = ref_pts_raw.sigmoid().unsqueeze(1) # [NQ, 1, 4] - - save_tensor(os.path.join(args.outdir, "ddec_query_embed"), query_embed) - save_tensor(os.path.join(args.outdir, "ddec_ref_pts_raw"), ref_pts_raw) - save_tensor(os.path.join(args.outdir, "ddec_ref_boxes_init"), reference_boxes) - - presence_token_w = ddec_sd["decoder.presence_token.weight"] # [1, D] - presence_out = presence_token_w[None].expand(1, 1, -1) # [1, 1, D] - save_tensor(os.path.join(args.outdir, "ddec_presence_token"), presence_out) - - enc_memory = fenc_output_bf.permute(1, 0, 2) # [5184, 1, D] seq-first - enc_pos = img_pe_flat.permute(1, 0, 2) # [5184, 1, D] seq-first - - valid_ratios = torch.ones(1, 1, 2, device=device) - spatial_shapes = torch.tensor([[H, H]], dtype=torch.long) - output = query_embed.unsqueeze(1).repeat(1, 1, 1) # [NQ, 1, D] - - coords_h = torch.arange(0, H, dtype=torch.float32) / H - coords_w = torch.arange(0, H, dtype=torch.float32) / H - - for layer_idx in range(6): - lp = f"decoder.layers.{layer_idx}." - - # Reference points input - reference_points_input = reference_boxes[:, :, None] * torch.cat([valid_ratios, valid_ratios], -1)[None, :] - query_sine_embed = gen_sineembed_for_position(reference_points_input[:, :, 0, :], D) - query_pos = mlp_forward(query_sine_embed, ddec_sd, "decoder.ref_point_head", 2) - - if layer_idx == 0: - save_tensor(os.path.join(args.outdir, "ddec_query_sine_0"), query_sine_embed) - save_tensor(os.path.join(args.outdir, "ddec_query_pos_0"), query_pos) - - # Box RPB - boxes_xyxy = box_cxcywh_to_xyxy(reference_boxes).permute(1, 0, 2) - bs_d, nq_d, _ = boxes_xyxy.shape - deltas_y = coords_h.view(1, -1, 1) - boxes_xyxy.reshape(-1, 1, 4)[:, :, 1:4:2] - deltas_y = deltas_y.view(bs_d, nq_d, -1, 2) - deltas_x = coords_w.view(1, -1, 1) - boxes_xyxy.reshape(-1, 1, 4)[:, :, 0:3:2] - deltas_x = deltas_x.view(bs_d, nq_d, -1, 2) - - deltas_x_log = deltas_x * 8 - deltas_x_log = torch.sign(deltas_x_log) * torch.log2(torch.abs(deltas_x_log) + 1.0) / np.log2(8) - deltas_y_log = deltas_y * 8 - deltas_y_log = torch.sign(deltas_y_log) * torch.log2(torch.abs(deltas_y_log) + 1.0) / np.log2(8) - deltas_x = deltas_x_log - deltas_y = deltas_y_log - - rpb_x = mlp_forward(deltas_x, ddec_sd, "decoder.boxRPB_embed_x", 2) - rpb_y = mlp_forward(deltas_y, ddec_sd, "decoder.boxRPB_embed_y", 2) - B_rpb = rpb_y.unsqueeze(3) + rpb_x.unsqueeze(2) - B_rpb = B_rpb.flatten(2, 3) - B_rpb = B_rpb.permute(0, 3, 1, 2).contiguous() - pres_rpb = torch.zeros(1, NH, 1, H*H, device=device) - memory_mask = torch.cat([pres_rpb, B_rpb], dim=2) - memory_mask = memory_mask.flatten(0, 1) - - if layer_idx == 0: - save_tensor(os.path.join(args.outdir, "ddec_rpb_mask_0"), memory_mask) - - # Self-attention - tgt_with_pres = torch.cat([presence_out, output], dim=0) - qpos_with_pres = torch.cat([torch.zeros_like(presence_out), query_pos], dim=0) - q_sa = k_sa = tgt_with_pres + qpos_with_pres - sa_w = ddec_sd[lp + "self_attn.in_proj_weight"] - sa_b = ddec_sd[lp + "self_attn.in_proj_bias"] - sa_ow = ddec_sd[lp + "self_attn.out_proj.weight"] - sa_ob = ddec_sd[lp + "self_attn.out_proj.bias"] - q_sa_bf = q_sa.permute(1, 0, 2) - k_sa_bf = k_sa.permute(1, 0, 2) - v_sa_bf = tgt_with_pres.permute(1, 0, 2) - sa_out = fused_mha_forward(q_sa_bf, k_sa_bf, v_sa_bf, sa_w, sa_b, sa_ow, sa_ob, NH) - tgt_with_pres = tgt_with_pres + sa_out.permute(1, 0, 2) - tgt_with_pres = F.layer_norm(tgt_with_pres, [D], - ddec_sd[lp + "norm2.weight"], ddec_sd[lp + "norm2.bias"]) - if layer_idx == 0: - save_tensor(os.path.join(args.outdir, "ddec_layer0_after_sa"), tgt_with_pres) - - # Text cross-attention - tgt_q_ca_text = tgt_with_pres + qpos_with_pres - ca_text_w = ddec_sd[lp + "ca_text.in_proj_weight"] - ca_text_b = ddec_sd[lp + "ca_text.in_proj_bias"] - ca_text_ow = ddec_sd[lp + "ca_text.out_proj.weight"] - ca_text_ob = ddec_sd[lp + "ca_text.out_proj.bias"] - w_q_ct, w_k_ct, w_v_ct = ca_text_w[:D], ca_text_w[D:2*D], ca_text_w[2*D:] - b_q_ct, b_k_ct, b_v_ct = ca_text_b[:D], ca_text_b[D:2*D], ca_text_b[2*D:] - tq_bf = tgt_q_ca_text.permute(1, 0, 2) - p_bf = prompt_sf.permute(1, 0, 2) - ct_out = multihead_attention_forward(tq_bf, p_bf, p_bf, - w_q_ct, b_q_ct, w_k_ct, b_k_ct, w_v_ct, b_v_ct, - ca_text_ow, ca_text_ob, NH, - key_padding_mask=txt_mask) - tgt_with_pres = tgt_with_pres + ct_out.permute(1, 0, 2) - tgt_with_pres = F.layer_norm(tgt_with_pres, [D], - ddec_sd[lp + "catext_norm.weight"], - ddec_sd[lp + "catext_norm.bias"]) - if layer_idx == 0: - save_tensor(os.path.join(args.outdir, "ddec_layer0_after_text_ca"), tgt_with_pres) - - # Image cross-attention with RPB - tgt_q_ca_img = tgt_with_pres + qpos_with_pres - mem_k = enc_memory + enc_pos - ca_w = ddec_sd[lp + "cross_attn.in_proj_weight"] - ca_b = ddec_sd[lp + "cross_attn.in_proj_bias"] - ca_ow = ddec_sd[lp + "cross_attn.out_proj.weight"] - ca_ob = ddec_sd[lp + "cross_attn.out_proj.bias"] - w_q_ci, w_k_ci, w_v_ci = ca_w[:D], ca_w[D:2*D], ca_w[2*D:] - b_q_ci, b_k_ci, b_v_ci = ca_b[:D], ca_b[D:2*D], ca_b[2*D:] - tqi_bf = tgt_q_ca_img.permute(1, 0, 2) - mk_bf = mem_k.permute(1, 0, 2) - mv_bf = enc_memory.permute(1, 0, 2) - ci_out = multihead_attention_forward(tqi_bf, mk_bf, mv_bf, - w_q_ci, b_q_ci, w_k_ci, b_k_ci, w_v_ci, b_v_ci, - ca_ow, ca_ob, NH, - attn_mask=memory_mask) - tgt_with_pres = tgt_with_pres + ci_out.permute(1, 0, 2) - tgt_with_pres = F.layer_norm(tgt_with_pres, [D], - ddec_sd[lp + "norm1.weight"], ddec_sd[lp + "norm1.bias"]) - if layer_idx == 0: - save_tensor(os.path.join(args.outdir, "ddec_layer0_after_img_ca"), tgt_with_pres) - - # FFN - ffn = F.linear(tgt_with_pres, ddec_sd[lp + "linear1.weight"], ddec_sd[lp + "linear1.bias"]) - ffn = F.relu(ffn) - ffn = F.linear(ffn, ddec_sd[lp + "linear2.weight"], ddec_sd[lp + "linear2.bias"]) - tgt_with_pres = tgt_with_pres + ffn - tgt_with_pres = F.layer_norm(tgt_with_pres, [D], - ddec_sd[lp + "norm3.weight"], ddec_sd[lp + "norm3.bias"]) - if layer_idx == 0: - save_tensor(os.path.join(args.outdir, "ddec_layer0_full_out"), tgt_with_pres) - - # Split presence and queries - presence_out = tgt_with_pres[:1] - output = tgt_with_pres[1:] - - # Box refinement - reference_before_sigmoid = inverse_sigmoid(reference_boxes) - normed_output = F.layer_norm(output, [D], - ddec_sd["decoder.norm.weight"], ddec_sd["decoder.norm.bias"]) - delta_unsig = mlp_forward(normed_output, ddec_sd, "decoder.bbox_embed", 3) - outputs_unsig = delta_unsig + reference_before_sigmoid - new_reference_points = outputs_unsig.sigmoid() - reference_boxes = new_reference_points.detach() - - save_tensor(os.path.join(args.outdir, f"ddec_layer{layer_idx}_out"), output) - save_tensor(os.path.join(args.outdir, f"ddec_layer{layer_idx}_refboxes"), reference_boxes) - if layer_idx == 0: - save_tensor(os.path.join(args.outdir, f"ddec_layer{layer_idx}_presence"), presence_out) - - # Final - normed_output = F.layer_norm(output, [D], - ddec_sd["decoder.norm.weight"], ddec_sd["decoder.norm.bias"]) - save_tensor(os.path.join(args.outdir, "ddec_normed_output"), normed_output) - save_tensor(os.path.join(args.outdir, "ddec_pred_boxes"), reference_boxes) - - pres_normed = F.layer_norm(presence_out, [D], - ddec_sd["decoder.presence_token_out_norm.weight"], - ddec_sd["decoder.presence_token_out_norm.bias"]) - pres_logit = mlp_forward(pres_normed, ddec_sd, "decoder.presence_token_head", 3) - save_tensor(os.path.join(args.outdir, "ddec_presence_logit"), pres_logit) - - # Scoring - scoring_prefix = "detector.dot_prod_scoring." - scoring_sd = {k[len(scoring_prefix):]: v for k, v in ckpt.items() - if k.startswith(scoring_prefix)} - - prompt_for_scoring = prompt_sf.clone() - orig = prompt_for_scoring.clone() - prompt_mlp_out = F.linear(prompt_for_scoring, scoring_sd["prompt_mlp.layers.0.weight"], - scoring_sd["prompt_mlp.layers.0.bias"]) - prompt_mlp_out = F.relu(prompt_mlp_out) - prompt_mlp_out = F.linear(prompt_mlp_out, scoring_sd["prompt_mlp.layers.1.weight"], - scoring_sd["prompt_mlp.layers.1.bias"]) - prompt_mlp_out = prompt_mlp_out + orig - prompt_mlp_out = F.layer_norm(prompt_mlp_out, [D], - scoring_sd["prompt_mlp.out_norm.weight"], - scoring_sd["prompt_mlp.out_norm.bias"]) - save_tensor(os.path.join(args.outdir, "scoring_prompt_mlp_out"), prompt_mlp_out) - - is_valid = (~txt_mask).float().permute(1, 0)[..., None] - num_valid = torch.clamp(is_valid.sum(dim=0), min=1.0) - pooled = (prompt_mlp_out * is_valid).sum(dim=0) / num_valid - save_tensor(os.path.join(args.outdir, "scoring_pooled"), pooled) - - proj_pooled = F.linear(pooled, scoring_sd["prompt_proj.weight"], - scoring_sd["prompt_proj.bias"]) - save_tensor(os.path.join(args.outdir, "scoring_proj_pooled"), proj_pooled) - - hs = normed_output.permute(1, 0, 2).unsqueeze(0) - proj_hs = F.linear(hs, scoring_sd["hs_proj.weight"], scoring_sd["hs_proj.bias"]) - save_tensor(os.path.join(args.outdir, "scoring_proj_hs"), proj_hs) - - scores = torch.matmul(proj_hs, proj_pooled.unsqueeze(-1)) - scores = scores * (1.0 / math.sqrt(D)) - scores = scores.clamp(-12.0, 12.0) - save_tensor(os.path.join(args.outdir, "scoring_class_scores"), scores) - - print("\nDone! All decoder tensors saved.") - - -if __name__ == "__main__": - main() diff --git a/tests/dump_fenc_from_package.py b/tests/dump_fenc_from_package.py deleted file mode 100644 index 4c62f6e..0000000 --- a/tests/dump_fenc_from_package.py +++ /dev/null @@ -1,469 +0,0 @@ -#!/usr/bin/env python3 -""" -Dump fusion encoder inputs and per-layer outputs from the ACTUAL SAM3 Python package. - -This script monkey-patches TransformerEncoder.forward to intercept the exact tensors -flowing through the fusion encoder during real inference. This ensures the reference -tensors are bit-identical to what the real model produces — no reimplementation. - -Usage: - cd ~/Documents/sam3 - uv run python /Users/pierre-antoine/Documents/sam3.cpp/tests/dump_fenc_from_package.py \ - [--image /path/to/image.jpg] [--text "yellow school bus"] \ - [--outdir /Users/pierre-antoine/Documents/sam3.cpp/tests/ref_fenc_pkg] -""" -import argparse -import os -import sys -import types - -# ── Mock triton for macOS (must happen before any torch/sam3 import) ── -def _install_triton_mock(): - """Install a mock triton package so SAM3 and torch._dynamo can import.""" - if "triton" in sys.modules: - return # Already real or already mocked - - triton = types.ModuleType("triton") - triton.__path__ = [] # Make it a package - triton.jit = lambda fn=None, **kw: fn if fn else (lambda f: f) - triton.autotune = lambda **kw: lambda f: f - triton.heuristics = lambda **kw: lambda f: f - triton.Config = type("Config", (), {"__init__": lambda self, **kw: None}) - - lang = types.ModuleType("triton.language") - lang.__path__ = [] - lang.constexpr = int - lang.dtype = type # torch._dynamo.utils needs this - triton.language = lang - - core = types.ModuleType("triton.language.core") - core.dtype = type - lang.core = core - - backends = types.ModuleType("triton.backends") - backends.__path__ = [] - compiler = types.ModuleType("triton.backends.compiler") - backends.compiler = compiler - - runtime = types.ModuleType("triton.runtime") - runtime.__path__ = [] - - triton_compiler = types.ModuleType("triton.compiler") - triton_compiler.__path__ = [] - triton_compiler_compiler = types.ModuleType("triton.compiler.compiler") - triton_compiler.compiler = triton_compiler_compiler - - # Wire submodule attributes (so triton.backends.compiler works) - triton.backends = backends - triton.runtime = runtime - triton.compiler = triton_compiler - lang.core = core - backends.compiler = compiler - triton_compiler.compiler = triton_compiler_compiler - - for name, mod in [ - ("triton", triton), - ("triton.language", lang), - ("triton.language.core", core), - ("triton.backends", backends), - ("triton.backends.compiler", compiler), - ("triton.runtime", runtime), - ("triton.compiler", triton_compiler), - ("triton.compiler.compiler", triton_compiler_compiler), - ]: - sys.modules[name] = mod - -_install_triton_mock() - -import numpy as np -import torch - -# Force CPU — SAM3 model builder uses device="cuda" in some places. -# Monkey-patch torch.zeros/torch.ones/torch.empty to replace device="cuda" with "cpu". -_orig_zeros = torch.zeros -_orig_ones = torch.ones -_orig_empty = torch.empty - -def _cpu_override(orig_fn): - def wrapper(*args, **kwargs): - dev = kwargs.get("device") - if isinstance(dev, str) and "cuda" in dev: - kwargs["device"] = "cpu" - elif isinstance(dev, torch.device) and dev.type == "cuda": - kwargs["device"] = "cpu" - return orig_fn(*args, **kwargs) - return wrapper - -_orig_arange = torch.arange -_orig_linspace = torch.linspace -_orig_full = torch.full -_orig_randn = torch.randn -_orig_rand = torch.rand -_orig_tensor = torch.tensor - -# Also neuter pin_memory — it tries to use MPS on macOS and causes device mismatch -_orig_pin_memory = torch.Tensor.pin_memory -torch.Tensor.pin_memory = lambda self, *a, **kw: self - -torch.zeros = _cpu_override(_orig_zeros) -torch.ones = _cpu_override(_orig_ones) -torch.empty = _cpu_override(_orig_empty) -torch.arange = _cpu_override(_orig_arange) -torch.linspace = _cpu_override(_orig_linspace) -torch.full = _cpu_override(_orig_full) -torch.randn = _cpu_override(_orig_randn) -torch.rand = _cpu_override(_orig_rand) -torch.tensor = _cpu_override(_orig_tensor) - -# --------------------------------------------------------------------------- -# Tensor save helpers (matching existing test infrastructure format) -# --------------------------------------------------------------------------- - -def save_tensor(path, t): - """Save tensor as .bin + .shape (PyTorch native layout).""" - t = t.detach().cpu().float().contiguous() - os.makedirs(os.path.dirname(path) or ".", exist_ok=True) - with open(path + ".bin", "wb") as f: - f.write(t.numpy().tobytes()) - with open(path + ".shape", "w") as f: - f.write(",".join(str(d) for d in t.shape)) - print(f" saved {path} shape={list(t.shape)} " - f"min={t.min().item():.4f} max={t.max().item():.4f}") - - -def save_tensor_i32(path, t): - t = t.detach().cpu().to(torch.int32).contiguous() - os.makedirs(os.path.dirname(path) or ".", exist_ok=True) - with open(path + ".bin", "wb") as f: - f.write(t.numpy().tobytes()) - with open(path + ".shape", "w") as f: - f.write(",".join(str(d) for d in t.shape)) - print(f" saved {path} shape={list(t.shape)} dtype=int32") - - -# --------------------------------------------------------------------------- -# Global capture dict — filled by the monkey-patched encoder forward -# --------------------------------------------------------------------------- - -captured = {} - - -def install_encoder_patch(): - """ - Monkey-patch TransformerEncoder.forward (the base class used by - TransformerEncoderFusion) to dump inputs and per-layer outputs. - - TransformerEncoderFusion.forward() does: - 1. Reshape src from seq-first to NCHW - 2. (Optional pooled text — disabled with add_pooled_text_to_img_feat=False) - 3. Call super().forward(src_NCHW, ..., prompt=prompt.transpose(0,1), ...) - - super().forward() (TransformerEncoder.forward) does: - 1. _prepare_multilevel_features → src_flatten [B, HW, D], lvl_pos_embed_flatten [B, HW, D] - 2. Loop through 6 encoder layers - - We patch TransformerEncoder.forward to capture the flattened tensors. - """ - from sam3.model.encoder import TransformerEncoder - from sam3.model.act_ckpt_utils import activation_ckpt_wrapper - - original_forward = TransformerEncoder.forward - - def patched_forward(self, src, src_key_padding_masks=None, pos=None, - prompt=None, prompt_key_padding_mask=None, - encoder_extra_kwargs=None): - # --- _prepare_multilevel_features (same as original) --- - ( - src_flatten, - key_padding_masks_flatten, - lvl_pos_embed_flatten, - level_start_index, - valid_ratios, - spatial_shapes, - ) = self._prepare_multilevel_features(src, src_key_padding_masks, pos) - - reference_points = self.get_reference_points( - spatial_shapes, valid_ratios, device=src_flatten.device - ) - - # ═══ DUMP INPUTS ═══ - # src_flatten: [B, HW, D] = [1, 5184, 256] batch-first image features - captured["fenc_input_tgt"] = src_flatten.detach().clone() - # lvl_pos_embed_flatten: [B, HW, D] = [1, 5184, 256] positional encoding - captured["fenc_input_pos"] = lvl_pos_embed_flatten.detach().clone() - # prompt: [B, T, D] = [1, 32, 256] batch-first text tokens - if prompt is not None: - captured["fenc_input_prompt"] = prompt.detach().clone() - # prompt_key_padding_mask: [B, T] = [1, 32] boolean (True=padding) - if prompt_key_padding_mask is not None: - captured["fenc_input_prompt_mask"] = prompt_key_padding_mask.detach().clone() - - # --- Layer loop (same as original) --- - output = src_flatten - for layer_idx, layer in enumerate(self.layers): - layer_kwargs = {} - - assert hasattr(layer, 'forward_pre') or hasattr(layer, 'forward_post') - layer_kwargs["memory"] = prompt - layer_kwargs["memory_key_padding_mask"] = prompt_key_padding_mask - layer_kwargs["query_pos"] = lvl_pos_embed_flatten - layer_kwargs["tgt"] = output - layer_kwargs["tgt_key_padding_mask"] = key_padding_masks_flatten - - if encoder_extra_kwargs is not None: - layer_kwargs.update(encoder_extra_kwargs) - output = activation_ckpt_wrapper(layer)( - **layer_kwargs, - act_ckpt_enable=self.training and self.use_act_checkpoint, - ) - - # ═══ DUMP PER-LAYER OUTPUT ═══ - captured[f"fenc_layer{layer_idx}_out"] = output.detach().clone() - - # --- Return (same as original) --- - return ( - output.transpose(0, 1), - ( - key_padding_masks_flatten.transpose(0, 1) - if key_padding_masks_flatten is not None - else None - ), - lvl_pos_embed_flatten.transpose(0, 1), - level_start_index, - spatial_shapes, - valid_ratios, - ) - - TransformerEncoder.forward = patched_forward - print("[PATCH] TransformerEncoder.forward patched for tensor dumping") - - -def install_text_encoder_hook(model): - """ - Hook into the text encoder output to capture text features and token info. - We intercept backbone.forward_text by wrapping it. - """ - original_forward_text = model.backbone.forward_text - - def patched_forward_text(text_list, **kwargs): - result = original_forward_text(text_list, **kwargs) - # language_features: [T, num_prompts, D] = [32, 1, 256] - if "language_features" in result: - captured["text_features_raw"] = result["language_features"].detach().clone() - if "language_mask" in result: - captured["text_mask_raw"] = result["language_mask"].detach().clone() - return result - - model.backbone.forward_text = patched_forward_text - print("[PATCH] backbone.forward_text wrapped for text feature capture") - - -def install_prompt_hook(model): - """ - Hook into _encode_prompt to capture the assembled prompt (text + geo + visual). - """ - original_encode_prompt = model._encode_prompt - - def patched_encode_prompt(*args, **kwargs): - prompt, prompt_mask, backbone_out = original_encode_prompt(*args, **kwargs) - # prompt: [seq, batch, dim] = [T+geo, 1, 256] seq-first - captured["prompt_assembled"] = prompt.detach().clone() - captured["prompt_mask_assembled"] = prompt_mask.detach().clone() - return prompt, prompt_mask, backbone_out - - model._encode_prompt = patched_encode_prompt - print("[PATCH] _encode_prompt wrapped for prompt capture") - - -def install_backbone_hook(model): - """ - Capture backbone FPN features after set_image. - """ - original_forward_image = model.backbone.forward_image - - def patched_forward_image(image): - result = original_forward_image(image) - if "backbone_fpn" in result: - for i, feat in enumerate(result["backbone_fpn"]): - captured[f"backbone_fpn_{i}"] = feat.detach().clone() - if "vision_pos_enc" in result: - for i, pe in enumerate(result["vision_pos_enc"]): - captured[f"vision_pos_enc_{i}"] = pe.detach().clone() - return result - - model.backbone.forward_image = patched_forward_image - print("[PATCH] backbone.forward_image wrapped for FPN capture") - - -# --------------------------------------------------------------------------- -# Main -# --------------------------------------------------------------------------- - -def main(): - parser = argparse.ArgumentParser(description="Dump fenc tensors from SAM3 package") - parser.add_argument("--image", default=None, - help="Test image path (default: tests/test_random.jpg)") - parser.add_argument("--text", default="yellow school bus", - help="Text prompt") - parser.add_argument("--outdir", - default="/Users/pierre-antoine/Documents/sam3.cpp/tests/ref_fenc_pkg", - help="Output directory for dumped tensors") - args = parser.parse_args() - - if args.image is None: - args.image = "/Users/pierre-antoine/Documents/sam3.cpp/tests/test_random.jpg" - - os.makedirs(args.outdir, exist_ok=True) - - # ── Load model ── - print("Loading SAM3 model...") - sys.path.insert(0, os.path.expanduser("~/Documents/sam3")) - - import sam3 - from sam3 import build_sam3_image_model - from sam3.model.sam3_image_processor import Sam3Processor - - sam3_pkg_dir = os.path.dirname(sam3.__file__) # ~/Documents/sam3/sam3/ - torch.set_default_dtype(torch.float32) - - bpe_path = os.path.join(sam3_pkg_dir, "assets", "bpe_simple_vocab_16e6.txt.gz") - if not os.path.exists(bpe_path): - # Try parent dir - bpe_path = os.path.join(sam3_pkg_dir, "..", "assets", "bpe_simple_vocab_16e6.txt.gz") - print(f" BPE path: {bpe_path} (exists={os.path.exists(bpe_path)})") - - ckpt_path = "/Users/pierre-antoine/Documents/sam3.cpp/raw_weights/sam3.pt" - print(f" Checkpoint: {ckpt_path} (exists={os.path.exists(ckpt_path)})") - - model = build_sam3_image_model( - bpe_path=bpe_path, - checkpoint_path=ckpt_path, - device="cpu", - load_from_HF=False, - ) - model = model.float() - model.eval() - - # ── Install all patches ── - install_encoder_patch() - install_text_encoder_hook(model) - install_prompt_hook(model) - install_backbone_hook(model) - - # ── Load and preprocess image ── - from PIL import Image - print(f"Loading image: {args.image}") - image = Image.open(args.image).convert("RGB") - - processor = Sam3Processor(model, confidence_threshold=0.5, device="cpu") - - # ── Step 1: Encode image (runs backbone) ── - print("\n=== Step 1: Encoding image ===") - state = processor.set_image(image) - - # Save backbone features - for key in sorted(captured.keys()): - if key.startswith("backbone_fpn_") or key.startswith("vision_pos_enc_"): - save_tensor(os.path.join(args.outdir, key), captured[key]) - - # ── Step 2: Run text prompt (runs text encoder + fusion encoder + decoder) ── - print(f"\n=== Step 2: Running text prompt '{args.text}' ===") - state = processor.set_text_prompt(prompt=args.text, state=state) - - # ── Save all captured tensors ── - print(f"\n=== Saving tensors to {args.outdir} ===") - - # Save prompt text - with open(os.path.join(args.outdir, "prompt.txt"), "w") as f: - f.write(args.text + "\n") - - # -- Fusion encoder inputs -- - print("\n-- Fusion encoder inputs --") - # fenc_input_tgt: [1, 5184, 256] batch-first image features - # Memory layout: 5184 blocks of 256 floats = same as ggml [256, 5184, 1] - if "fenc_input_tgt" in captured: - save_tensor(os.path.join(args.outdir, "fenc_input_tgt"), captured["fenc_input_tgt"]) - - # fenc_input_pos: [1, 5184, 256] batch-first positional encoding - if "fenc_input_pos" in captured: - save_tensor(os.path.join(args.outdir, "fenc_input_pos"), captured["fenc_input_pos"]) - - # fenc_input_prompt: [1, T, 256] batch-first prompt tokens - if "fenc_input_prompt" in captured: - save_tensor(os.path.join(args.outdir, "fenc_input_prompt"), captured["fenc_input_prompt"]) - - # fenc_input_prompt_mask: [1, T] boolean (True=padding) - if "fenc_input_prompt_mask" in captured: - mask = captured["fenc_input_prompt_mask"] - save_tensor(os.path.join(args.outdir, "fenc_input_prompt_mask"), - mask.float()) - # Also save as attention bias (0.0 for valid, -1e9 for padding) - # This is what the C++ uses directly - attn_bias = torch.where(mask, torch.tensor(-1e9), torch.tensor(0.0)) - save_tensor(os.path.join(args.outdir, "fenc_attn_bias"), attn_bias) - - # -- Also save in format compatible with existing phase 5 test -- - print("\n-- Phase 5 compatible format --") - if "fenc_input_tgt" in captured: - # [1, 5184, 256] → same binary layout as [32, 1, 256] seq-first - save_tensor(os.path.join(args.outdir, "fenc_img_input"), captured["fenc_input_tgt"]) - if "fenc_input_pos" in captured: - save_tensor(os.path.join(args.outdir, "fenc_pos_embed"), captured["fenc_input_pos"]) - - # Text features in seq-first format [T, B, D] for compatibility - if "prompt_assembled" in captured: - # prompt_assembled is [T, B, D] seq-first - save_tensor(os.path.join(args.outdir, "text_features"), captured["prompt_assembled"]) - - # -- Fusion encoder per-layer outputs -- - print("\n-- Fusion encoder per-layer outputs --") - for i in range(6): - key = f"fenc_layer{i}_out" - if key in captured: - save_tensor(os.path.join(args.outdir, key), captured[key]) - - # -- Text encoder raw output -- - print("\n-- Text encoder output --") - if "text_features_raw" in captured: - save_tensor(os.path.join(args.outdir, "text_features_raw"), captured["text_features_raw"]) - if "text_mask_raw" in captured: - save_tensor(os.path.join(args.outdir, "text_mask_raw"), captured["text_mask_raw"].float()) - - # -- Also save neck features for potential full pipeline test -- - print("\n-- Neck detector features --") - # The detector uses backbone_fpn levels. Save the last 3 for the seg head. - for i in range(4): - key = f"backbone_fpn_{i}" - if key in captured: - t = captured[key] - # Save as neck_det_{i} in NCHW format - save_tensor(os.path.join(args.outdir, f"neck_det_{i}"), t) - - # -- Summary -- - print(f"\n=== Summary ===") - print(f"Total tensors captured: {len(captured)}") - if "fenc_input_tgt" in captured: - print(f" fenc image features: {list(captured['fenc_input_tgt'].shape)}") - if "fenc_input_prompt" in captured: - t = captured["fenc_input_prompt"] - print(f" fenc prompt tokens: {list(t.shape)}") - if "fenc_input_prompt_mask" in captured: - mask = captured["fenc_input_prompt_mask"] - n_valid = (~mask).sum().item() - n_total = mask.numel() - print(f" prompt mask: {n_valid}/{n_total} valid tokens") - if "fenc_layer5_out" in captured: - print(f" fenc output: {list(captured['fenc_layer5_out'].shape)}") - - # Check for detection results - if "boxes" in state and len(state.get("boxes", [])) > 0: - print(f"\n Detections: {len(state['boxes'])} objects found") - else: - print(f"\n Detections: none (may need different prompt/image)") - - print(f"\nAll tensors saved to {args.outdir}/") - print("Next: run the C++ test_debug_fenc against these references.") - - -if __name__ == "__main__": - main() diff --git a/tests/dump_geom_encoder_reference.py b/tests/dump_geom_encoder_reference.py deleted file mode 100644 index 4b2d37f..0000000 --- a/tests/dump_geom_encoder_reference.py +++ /dev/null @@ -1,482 +0,0 @@ -#!/usr/bin/env python3 -""" -Dump geometry encoder reference tensors for numerical comparison against sam3.cpp. - -Tests the SequenceGeometryEncoder with exemplar box prompts. -Dumps intermediate tensors at each stage: - 1. Input box coordinates (CxCyWH normalized) - 2. Direct box projection output - 3. ROI align pooled features - 4. Box positional encoding - 5. Label embedding - 6. Combined box embedding - 7. CLS token - 8. Final projection output - 9. Per-layer transformer outputs - 10. Final output - -Also tests the dummy prompt case (no exemplars, just CLS token). - -Usage: - cd ~/Documents/sam3 - uv run python /Users/pierre-antoine/Documents/sam3.cpp/tests/dump_geom_encoder_reference.py \ - --checkpoint raw_weights/sam3.pt \ - --prephase-ref /Users/pierre-antoine/Documents/sam3.cpp/tests/ref_phase3 \ - --outdir /Users/pierre-antoine/Documents/sam3.cpp/tests/ref_geom -""" - -import argparse -import math -import os -import sys -from typing import Dict, List, Tuple - -import numpy as np -import torch -import torch.nn.functional as F -import torchvision - -# ── Tensor save utilities ─────────────────────────────────────────────────── - -def save_raw(path: str, arr: np.ndarray, shape) -> None: - os.makedirs(os.path.dirname(path) or ".", exist_ok=True) - with open(path + ".bin", "wb") as f: - f.write(arr.astype(np.float32, copy=False).tobytes()) - with open(path + ".shape", "w", encoding="utf-8") as f: - f.write(",".join(str(int(d)) for d in shape)) - - -def save_ggml_bnd(path: str, x_bnd: torch.Tensor) -> None: - """Save [B, N, D] as ggml [D, N, B].""" - assert x_bnd.ndim == 3 - x = x_bnd.detach().cpu().float().contiguous() - b, n, d = x.shape - save_raw(path, x.numpy(), (d, n, b)) - - -def save_ggml_sbd(path: str, x_sbd: torch.Tensor) -> None: - """Save [S, B, D] (seq-first) as ggml [D, S, B].""" - assert x_sbd.ndim == 3 - x = x_sbd.detach().cpu().float().contiguous() - s, b, d = x.shape - # Permute to [B, S, D] then save as ggml [D, S, B] - x_bsd = x.permute(1, 0, 2).contiguous() - save_raw(path, x_bsd.numpy(), (d, s, b)) - - -def save_ggml_nchw(path: str, x_nchw: torch.Tensor) -> None: - """Save [N, C, H, W] as ggml [C, W, H, N].""" - assert x_nchw.ndim == 4 - x = x_nchw.detach().cpu().float().contiguous() - n, c, h, w = x.shape - x_nhwc = x.permute(0, 2, 3, 1).contiguous() - save_raw(path, x_nhwc.numpy(), (c, w, h, n)) - - -def save_ggml_1d(path: str, x: torch.Tensor) -> None: - """Save a 1D tensor.""" - x = x.detach().cpu().float().contiguous() - save_raw(path, x.numpy(), tuple(x.shape)) - - -def load_tensor(path: str) -> torch.Tensor: - with open(path + ".shape", "r") as f: - shape = [int(x) for x in f.read().strip().split(",") if x] - data = np.fromfile(path + ".bin", dtype=np.float32).reshape(shape) - return torch.from_numpy(data) - - -# ── Sinusoidal position encoding (matches PositionEmbeddingSine) ────────── - -def sine_encode_xy(x: torch.Tensor, y: torch.Tensor, num_pos_feats: int = 128, - temperature: int = 10000, scale: float = 2 * math.pi): - """Encode (x, y) with sinusoidal PE. x, y: [N] normalized to [0,1].""" - x_embed = x * scale - y_embed = y * scale - - dim_t = torch.arange(num_pos_feats, dtype=torch.float32, device=x.device) - dim_t = temperature ** (2 * (dim_t // 2) / num_pos_feats) - - pos_x = x_embed[:, None] / dim_t # [N, 128] - pos_y = y_embed[:, None] / dim_t # [N, 128] - pos_x = torch.stack((pos_x[:, 0::2].sin(), pos_x[:, 1::2].cos()), dim=2).flatten(1) # [N, 128] - pos_y = torch.stack((pos_y[:, 0::2].sin(), pos_y[:, 1::2].cos()), dim=2).flatten(1) # [N, 128] - return pos_x, pos_y - - -def sine_encode_boxes(cx, cy, w, h, num_pos_feats=128, temperature=10000): - """Encode box (cx, cy, w, h) with sinusoidal PE. Returns [N, 258].""" - pos_x, pos_y = sine_encode_xy(cx, cy, num_pos_feats, temperature) - pos = torch.cat((pos_y, pos_x, h[:, None], w[:, None]), dim=1) # [N, 258] - return pos - - -# ── Box coordinate conversion ────────────────────────────────────────────── - -def box_cxcywh_to_xyxy(x): - """Convert CxCyWH to XYXY format.""" - cx, cy, w, h = x.unbind(-1) - x0 = cx - 0.5 * w - y0 = cy - 0.5 * h - x1 = cx + 0.5 * w - y1 = cy + 0.5 * h - return torch.stack([x0, y0, x1, y1], dim=-1) - - -# ── ROI Align implementation ──────────────────────────────────────────────── - -def do_roi_align(img_feats_nchw, boxes_cxcywh, roi_size=7): - """ - ROI Align from image features using box coordinates in CxCyWH [0,1] format. - - Args: - img_feats_nchw: [B, C, H, W] backbone feature map - boxes_cxcywh: [N_boxes, B, 4] in normalized CxCyWH format - roi_size: output spatial size - - Returns: - roi_pooled: [N_boxes*B, C, roi_size, roi_size] - """ - n_boxes, bs, _ = boxes_cxcywh.shape - H, W = img_feats_nchw.shape[-2:] - - boxes_xyxy = box_cxcywh_to_xyxy(boxes_cxcywh) - scale = torch.tensor([W, H, W, H], dtype=boxes_xyxy.dtype, device=boxes_xyxy.device) - scale = scale.view(1, 1, 4) - boxes_xyxy = boxes_xyxy * scale - - # roi_align expects list of [N_boxes, 4] per batch element - rois_list = boxes_xyxy.float().transpose(0, 1).unbind(0) - sampled = torchvision.ops.roi_align(img_feats_nchw, rois_list, roi_size) - return sampled # [B*N_boxes, C, roi_size, roi_size] - - -# ── Main dump logic ─────────────────────────────────────────────────────── - -def dump_geometry_encoder( - ckpt: Dict[str, torch.Tensor], - img_feats_hwc: torch.Tensor, # [H*W, B, C] seq-first image features - img_feats_nchw: torch.Tensor, # [B, C, H, W] image features for ROI align - img_pe_hwc: torch.Tensor, # [H*W, B, C] image positional encoding - boxes_cxcywh: torch.Tensor, # [N_boxes, B, 4] normalized CxCyWH [0,1] - box_labels: torch.Tensor, # [N_boxes, B] long (0=pos, 1=neg) - box_mask: torch.Tensor, # [B, N_boxes] bool (True=padded) - outdir: str, - case_name: str, -): - D = 256 - roi_size = 7 - case_dir = os.path.join(outdir, case_name) - os.makedirs(case_dir, exist_ok=True) - - n_boxes, bs = boxes_cxcywh.shape[:2] - print(f"\n=== Geometry Encoder: {case_name} ===") - print(f" boxes: {n_boxes}, batch: {bs}, img_feats: {img_feats_hwc.shape}") - - # Save input coordinates - save_raw(os.path.join(case_dir, "input_boxes_cxcywh"), - boxes_cxcywh.detach().cpu().float().numpy(), - boxes_cxcywh.shape) - save_raw(os.path.join(case_dir, "input_box_labels"), - box_labels.detach().cpu().float().numpy(), - box_labels.shape) - - # ── 1. Direct box projection: Linear(4, 256) ─────────────────────── - box_proj_w = ckpt["detector.geometry_encoder.boxes_direct_project.weight"].float() - box_proj_b = ckpt["detector.geometry_encoder.boxes_direct_project.bias"].float() - boxes_direct = F.linear(boxes_cxcywh, box_proj_w, box_proj_b) # [N, B, D] - save_ggml_sbd(os.path.join(case_dir, "boxes_direct_proj"), boxes_direct) - print(f" boxes_direct_proj: {boxes_direct.shape} " - f"mean={boxes_direct.mean():.6f} std={boxes_direct.std():.6f}") - - # ── 2. ROI Align pooled features ──────────────────────────────────── - # img_pre_norm on image features before pooling - img_pre_norm_w = ckpt["detector.geometry_encoder.img_pre_norm.weight"].float() - img_pre_norm_b = ckpt["detector.geometry_encoder.img_pre_norm.bias"].float() - # Apply LayerNorm to seq-first features for pooling (operate on the NCHW version) - img_normed_hwc = F.layer_norm(img_feats_hwc, [D], img_pre_norm_w, img_pre_norm_b) - H, W = img_feats_nchw.shape[-2:] - img_normed_nchw = img_normed_hwc.permute(1, 2, 0).reshape(bs, D, H, W) - - if n_boxes > 0: - roi_pooled = do_roi_align(img_normed_nchw, boxes_cxcywh, roi_size) - # roi_pooled: [B*N_boxes, D, roi_size, roi_size] - - # boxes_pool_project is Conv2d(D, D, roi_size) - pool_proj_w = ckpt["detector.geometry_encoder.boxes_pool_project.weight"].float() - pool_proj_b = ckpt["detector.geometry_encoder.boxes_pool_project.bias"].float() - pool_proj = F.conv2d(roi_pooled, pool_proj_w, pool_proj_b) # [B*N, D, 1, 1] - pool_proj = pool_proj.view(bs, n_boxes, D).transpose(0, 1) # [N, B, D] - save_ggml_sbd(os.path.join(case_dir, "boxes_pool_proj"), pool_proj) - print(f" boxes_pool_proj: {pool_proj.shape} " - f"mean={pool_proj.mean():.6f} std={pool_proj.std():.6f}") - else: - pool_proj = torch.zeros(0, bs, D) - - # ── 3. Box positional encoding: sinusoidal PE + Linear(258, 256) ─── - if n_boxes > 0: - cx = boxes_cxcywh[:, :, 0].flatten() - cy = boxes_cxcywh[:, :, 1].flatten() - w = boxes_cxcywh[:, :, 2].flatten() - h = boxes_cxcywh[:, :, 3].flatten() - pos_enc = sine_encode_boxes(cx, cy, w, h) # [N*B, 258] - pos_enc = pos_enc.view(n_boxes, bs, 258) # [N, B, 258] - save_ggml_sbd(os.path.join(case_dir, "boxes_pos_enc_raw"), pos_enc) - - pos_proj_w = ckpt["detector.geometry_encoder.boxes_pos_enc_project.weight"].float() - pos_proj_b = ckpt["detector.geometry_encoder.boxes_pos_enc_project.bias"].float() - pos_proj = F.linear(pos_enc, pos_proj_w, pos_proj_b) # [N, B, D] - save_ggml_sbd(os.path.join(case_dir, "boxes_pos_proj"), pos_proj) - print(f" boxes_pos_proj: {pos_proj.shape} " - f"mean={pos_proj.mean():.6f} std={pos_proj.std():.6f}") - else: - pos_proj = torch.zeros(0, bs, D) - - # ── 4. Label embedding ─────────────────────────────────────────────── - label_embed_w = ckpt["detector.geometry_encoder.label_embed.weight"].float() # [2, D] - if n_boxes > 0: - type_embed = F.embedding(box_labels.long(), label_embed_w) # [N, B, D] - save_ggml_sbd(os.path.join(case_dir, "label_embed"), type_embed) - print(f" label_embed: {type_embed.shape}") - else: - type_embed = torch.zeros(0, bs, D) - - # ── 5. Combined box embedding ──────────────────────────────────────── - if n_boxes > 0: - boxes_embed = boxes_direct + pool_proj + pos_proj + type_embed # [N, B, D] - else: - boxes_embed = torch.zeros(0, bs, D) - save_ggml_sbd(os.path.join(case_dir, "boxes_combined_embed"), boxes_embed) - print(f" boxes_combined: {boxes_embed.shape}") - - # ── 6. CLS token ──────────────────────────────────────────────────── - cls_embed_w = ckpt["detector.geometry_encoder.cls_embed.weight"].float() # [1, D] - cls = cls_embed_w.view(1, 1, D).repeat(1, bs, 1) # [1, B, D] - cls_mask = torch.zeros(bs, 1, dtype=torch.bool) - - # Concatenate: [boxes, CLS] - # Inline concat_padded_sequences for simplicity (no mask padding needed here) - if n_boxes > 0: - final_embeds = torch.cat([boxes_embed, cls], dim=0) # [N+1, B, D] - final_mask = torch.cat([box_mask, cls_mask], dim=1) # [B, N+1] - else: - final_embeds = cls # [1, B, D] - final_mask = cls_mask # [B, 1] - save_ggml_sbd(os.path.join(case_dir, "pre_final_proj"), final_embeds) - print(f" pre_final_proj: {final_embeds.shape}") - - # ── 7. Final projection + norm ────────────────────────────────────── - final_proj_w = ckpt["detector.geometry_encoder.final_proj.weight"].float() - final_proj_b = ckpt["detector.geometry_encoder.final_proj.bias"].float() - norm_w = ckpt["detector.geometry_encoder.norm.weight"].float() - norm_b = ckpt["detector.geometry_encoder.norm.bias"].float() - - final_embeds = F.layer_norm( - F.linear(final_embeds, final_proj_w, final_proj_b), - [D], norm_w, norm_b - ) - save_ggml_sbd(os.path.join(case_dir, "post_final_proj"), final_embeds) - print(f" post_final_proj: {final_embeds.shape}") - - # ── 8. Transformer layers (3 layers) ───────────────────────────────── - for layer_idx in range(3): - lp = f"detector.geometry_encoder.encode.{layer_idx}." - - # Self-attention (pre-norm, pos_enc_at_attn=False) - shortcut = final_embeds - xn = F.layer_norm(final_embeds, [D], - ckpt[lp + "norm1.weight"].float(), - ckpt[lp + "norm1.bias"].float()) - - # Self-attention with fused in_proj - sa_w = ckpt[lp + "self_attn.in_proj_weight"].float() # [3D, D] - sa_b = ckpt[lp + "self_attn.in_proj_bias"].float() # [3D] - sa_ow = ckpt[lp + "self_attn.out_proj.weight"].float() - sa_ob = ckpt[lp + "self_attn.out_proj.bias"].float() - - # pos_enc_at_attn=False, so Q=K=V=xn (no pos added) - # key_padding_mask = final_mask - S, B, _ = xn.shape - qkv = F.linear(xn, sa_w, sa_b) # [S, B, 3D] - q, k, v = qkv.chunk(3, dim=-1) - n_heads = 8 - hd = D // n_heads - - q = q.reshape(S, B, n_heads, hd).permute(1, 2, 0, 3) - k = k.reshape(S, B, n_heads, hd).permute(1, 2, 0, 3) - v = v.reshape(S, B, n_heads, hd).permute(1, 2, 0, 3) - - # Apply key_padding_mask - attn_mask = None - if final_mask.any(): - # final_mask: [B, S] True=padded -> attn_mask: [B, 1, 1, S] with -inf - attn_mask = final_mask.unsqueeze(1).unsqueeze(2).float() * (-1e9) - - sa_out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask) - sa_out = sa_out.permute(2, 0, 1, 3).reshape(S, B, D) - sa_out = F.linear(sa_out, sa_ow, sa_ob) - final_embeds = shortcut + sa_out - - save_ggml_sbd(os.path.join(case_dir, f"layer{layer_idx}_after_sa"), final_embeds) - - # Cross-attention (pre-norm, pos_enc_at_cross_attn_keys=True) - shortcut = final_embeds - xn = F.layer_norm(final_embeds, [D], - ckpt[lp + "norm2.weight"].float(), - ckpt[lp + "norm2.bias"].float()) - - ca_w = ckpt[lp + "cross_attn_image.in_proj_weight"].float() # [3D, D] - ca_b = ckpt[lp + "cross_attn_image.in_proj_bias"].float() - ca_ow = ckpt[lp + "cross_attn_image.out_proj.weight"].float() - ca_ob = ckpt[lp + "cross_attn_image.out_proj.bias"].float() - - # Q from xn (no pos), K from img_feats + img_pe, V from img_feats - q_w, k_w, v_w = ca_w.chunk(3, dim=0) - q_b, k_b, v_b = ca_b.chunk(3, dim=0) - - # pos_enc_at_cross_attn_queries=False, so Q = proj(xn) - # pos_enc_at_cross_attn_keys=True, so K = proj(img_feats + img_pe) - q_proj = F.linear(xn, q_w, q_b) # [S_q, B, D] - k_input = img_feats_hwc + img_pe_hwc # [S_kv, B, D] - k_proj = F.linear(k_input, k_w, k_b) # [S_kv, B, D] - v_proj = F.linear(img_feats_hwc, v_w, v_b) # [S_kv, B, D] - - S_q = q_proj.shape[0] - S_kv = k_proj.shape[0] - q_proj = q_proj.reshape(S_q, B, n_heads, hd).permute(1, 2, 0, 3) - k_proj = k_proj.reshape(S_kv, B, n_heads, hd).permute(1, 2, 0, 3) - v_proj = v_proj.reshape(S_kv, B, n_heads, hd).permute(1, 2, 0, 3) - - ca_out = F.scaled_dot_product_attention(q_proj, k_proj, v_proj) - ca_out = ca_out.permute(2, 0, 1, 3).reshape(S_q, B, D) - ca_out = F.linear(ca_out, ca_ow, ca_ob) - final_embeds = shortcut + ca_out - - save_ggml_sbd(os.path.join(case_dir, f"layer{layer_idx}_after_ca"), final_embeds) - - # FFN (pre-norm, ReLU) - shortcut = final_embeds - xn = F.layer_norm(final_embeds, [D], - ckpt[lp + "norm3.weight"].float(), - ckpt[lp + "norm3.bias"].float()) - ffn = F.linear(xn, ckpt[lp + "linear1.weight"].float(), - ckpt[lp + "linear1.bias"].float()) - ffn = F.relu(ffn) - ffn = F.linear(ffn, ckpt[lp + "linear2.weight"].float(), - ckpt[lp + "linear2.bias"].float()) - final_embeds = shortcut + ffn - - save_ggml_sbd(os.path.join(case_dir, f"layer{layer_idx}_after_ffn"), final_embeds) - print(f" layer{layer_idx}: shape={final_embeds.shape} " - f"mean={final_embeds.mean():.6f} std={final_embeds.std():.6f}") - - # ── 9. Final encode norm ──────────────────────────────────────────── - enc_norm_w = ckpt["detector.geometry_encoder.encode_norm.weight"].float() - enc_norm_b = ckpt["detector.geometry_encoder.encode_norm.bias"].float() - final_embeds = F.layer_norm(final_embeds, [D], enc_norm_w, enc_norm_b) - save_ggml_sbd(os.path.join(case_dir, "geom_output"), final_embeds) - print(f" geom_output: {final_embeds.shape} " - f"mean={final_embeds.mean():.6f} std={final_embeds.std():.6f}") - - # Also save the mask - save_raw(os.path.join(case_dir, "geom_mask"), - final_mask.cpu().float().numpy(), final_mask.shape) - - return final_embeds, final_mask - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--checkpoint", required=True) - parser.add_argument("--prephase-ref", required=True, - help="Phase 3 reference dir with neck features") - parser.add_argument("--outdir", required=True) - args = parser.parse_args() - - D = 256 - H = 72 - - # Load checkpoint - print("Loading checkpoint...") - ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=True) - - # Load backbone features from Phase 3 reference (saved in NCHW format) - print("Loading Phase 3 reference features...") - neck_det_2 = load_tensor(os.path.join(args.prephase_ref, "neck_det_2")) # [1, 256, 72, 72] - print(f" neck_det_2: {neck_det_2.shape}") - - # Image features in sequence-first format [H*W, B, C] - img_feats_hwc = neck_det_2.flatten(2).permute(2, 0, 1) # [H*W, B, C] - - # Image positional encoding [H*W, B, C] — sinusoidal PE computed on the fly - def compute_sine_pe(h, w, num_pos_feats=128, temperature=10000, scale=2*math.pi): - """Sinusoidal PE matching SAM3's PositionEmbeddingSine.""" - not_mask = torch.ones(1, h, w, dtype=torch.float32) - y_embed = not_mask.cumsum(1, dtype=torch.float32) - x_embed = not_mask.cumsum(2, dtype=torch.float32) - eps = 1e-6 - y_embed = (y_embed - 0.5) / (y_embed[:, -1:, :] + eps) * scale - x_embed = (x_embed - 0.5) / (x_embed[:, :, -1:] + eps) * scale - - dim_t = torch.arange(num_pos_feats, dtype=torch.float32) - dim_t = temperature ** (2 * (dim_t // 2) / num_pos_feats) - - pos_x = x_embed[:, :, :, None] / dim_t - pos_y = y_embed[:, :, :, None] / dim_t - pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3) - pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3) - pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) # [1, D, H, W] - return pos - - img_pe_nchw = compute_sine_pe(H, H, num_pos_feats=D // 2) # [1, 256, 72, 72] - img_pe_hwc = img_pe_nchw.flatten(2).permute(2, 0, 1) # [H*W, 1, C] - print(f" img_pe: {img_pe_hwc.shape}") - - # ══════════════════════════════════════════════════════════════════════ - # Test Case 1: Dummy prompt (no exemplars, just CLS) - # ══════════════════════════════════════════════════════════════════════ - dummy_boxes = torch.zeros(0, 1, 4) - dummy_labels = torch.zeros(0, 1, dtype=torch.long) - dummy_mask = torch.zeros(1, 0, dtype=torch.bool) - - dump_geometry_encoder( - ckpt, img_feats_hwc, neck_det_2, img_pe_hwc, - dummy_boxes, dummy_labels, dummy_mask, - args.outdir, "dummy_prompt" - ) - - # ══════════════════════════════════════════════════════════════════════ - # Test Case 2: Single positive exemplar box - # Box at center of image: cx=0.5, cy=0.5, w=0.3, h=0.3 - # ══════════════════════════════════════════════════════════════════════ - single_box = torch.tensor([[[0.5, 0.5, 0.3, 0.3]]]) # [1, 1, 4] - single_label = torch.tensor([[0]], dtype=torch.long) # [1, 1] positive=0 - single_mask = torch.tensor([[False]]) # [1, 1] - - dump_geometry_encoder( - ckpt, img_feats_hwc, neck_det_2, img_pe_hwc, - single_box, single_label, single_mask, - args.outdir, "single_box" - ) - - # ══════════════════════════════════════════════════════════════════════ - # Test Case 3: Two exemplar boxes (one positive, one negative) - # ══════════════════════════════════════════════════════════════════════ - two_boxes = torch.tensor([ - [[0.3, 0.4, 0.2, 0.25]], # box 1: positive - [[0.7, 0.6, 0.15, 0.2]], # box 2: negative - ]) # [2, 1, 4] - two_labels = torch.tensor([[0], [1]], dtype=torch.long) # [2, 1] - two_mask = torch.tensor([[False, False]]) # [1, 2] - - dump_geometry_encoder( - ckpt, img_feats_hwc, neck_det_2, img_pe_hwc, - two_boxes, two_labels, two_mask, - args.outdir, "two_boxes" - ) - - print(f"\nAll reference tensors saved to {args.outdir}/") - - -if __name__ == "__main__": - main() diff --git a/tests/dump_phase3_reference.py b/tests/dump_phase3_reference.py deleted file mode 100644 index c5197c9..0000000 --- a/tests/dump_phase3_reference.py +++ /dev/null @@ -1,427 +0,0 @@ -#!/usr/bin/env python3 -""" -Phase 3 Numerical Audit — dump ALL intermediate tensors from the ViT backbone, -neck (SimpleFPN), sinusoidal PE, and RoPE for numerical comparison against C++. - -Usage: - uv run python tests/dump_phase3_reference.py \ - --checkpoint raw_weights/sam3.pt \ - --image tests/test_random.jpg \ - --outdir tests/ref_phase3 -""" -import argparse -import math -import os -import sys -from functools import partial -from typing import Tuple - -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F -from PIL import Image -from torchvision.transforms import v2 - - -# ── Helpers ──────────────────────────────────────────────────────────────── - -def save_tensor(path, t): - t = t.detach().cpu().float().contiguous() - os.makedirs(os.path.dirname(path) or ".", exist_ok=True) - with open(path + ".bin", "wb") as f: - f.write(t.numpy().tobytes()) - with open(path + ".shape", "w") as f: - f.write(",".join(str(d) for d in t.shape)) - print(f" saved {path} shape={list(t.shape)} " - f"min={t.min().item():.6f} max={t.max().item():.6f} mean={t.mean().item():.6f}") - - -# ── Inlined from vitdet.py ───────────────────────────────────────────────── - -def init_t_xy(end_x, end_y, scale=1.0, offset=0): - t = torch.arange(end_x * end_y, dtype=torch.float32) - t_x = (t % end_x).float() - t_y = torch.div(t, end_x, rounding_mode="floor").float() - return t_x * scale + offset, t_y * scale + offset - - -def compute_axial_cis(dim, end_x, end_y, theta=10000.0, scale_pos=1.0, offset=0): - freqs_x = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim)) - freqs_y = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim)) - t_x, t_y = init_t_xy(end_x, end_y, scale_pos, offset) - freqs_x = torch.outer(t_x, freqs_x) - freqs_y = torch.outer(t_y, freqs_y) - freqs_cis_x = torch.polar(torch.ones_like(freqs_x), freqs_x) - freqs_cis_y = torch.polar(torch.ones_like(freqs_y), freqs_y) - return torch.cat([freqs_cis_x, freqs_cis_y], dim=-1) - - -def reshape_for_broadcast(freqs_cis, x): - ndim = x.ndim - shape = [d if i >= ndim - 2 else 1 for i, d in enumerate(x.shape)] - return freqs_cis.view(*shape) - - -def apply_rotary_enc(xq, xk, freqs_cis, repeat_freqs_k=False): - xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) - xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2)) if xk.shape[-2] != 0 else None - freqs_cis = reshape_for_broadcast(freqs_cis, xq_) - xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3) - if xk_ is None: - return xq_out.type_as(xq), xk - if repeat_freqs_k: - r = xk_.shape[-2] // xq_.shape[-2] - freqs_cis = freqs_cis.repeat(*([1] * (freqs_cis.ndim - 2)), r, 1) - xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3) - return xq_out.type_as(xq), xk_out.type_as(xk) - - -def window_partition(x, window_size): - B, H, W, C = x.shape - pad_h = (window_size - H % window_size) % window_size - pad_w = (window_size - W % window_size) % window_size - if pad_h > 0 or pad_w > 0: - x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h)) - Hp, Wp = H + pad_h, W + pad_w - x = x.view(B, Hp // window_size, window_size, Wp // window_size, window_size, C) - windows = x.permute(0, 1, 3, 2, 4, 5).reshape(-1, window_size, window_size, C) - return windows, (Hp, Wp) - - -def window_unpartition(windows, window_size, pad_hw, hw): - Hp, Wp = pad_hw - H, W = hw - B = windows.shape[0] // (Hp * Wp // window_size // window_size) - x = windows.reshape(B, Hp // window_size, Wp // window_size, window_size, window_size, -1) - x = x.permute(0, 1, 3, 2, 4, 5).reshape(B, Hp, Wp, -1) - if Hp > H or Wp > W: - x = x[:, :H, :W, :] - return x - - -def get_abs_pos(abs_pos, has_cls_token, hw, tiling=False): - h, w = hw - if has_cls_token: - abs_pos = abs_pos[:, 1:] - xy_num = abs_pos.shape[1] - size = int(math.sqrt(xy_num)) - if size != h or size != w: - new_abs_pos = abs_pos.reshape(1, size, size, -1).permute(0, 3, 1, 2) - if tiling: - new_abs_pos = new_abs_pos.tile( - [1, 1] + [x // y + 1 for x, y in zip((h, w), new_abs_pos.shape[2:])] - )[:, :, :h, :w] - else: - new_abs_pos = F.interpolate(new_abs_pos, size=(h, w), mode="bicubic", align_corners=False) - return new_abs_pos.permute(0, 2, 3, 1) - return abs_pos.reshape(1, h, w, -1) - - -# ── Sinusoidal PE (from position_encoding.py) ──────────────────────────── - -def sinusoidal_pe_2d(H, W, num_pos_feats=256, temperature=10000, scale=2*math.pi): - """Match PositionEmbeddingSine.forward() exactly.""" - half = num_pos_feats // 2 # 128 - - y_embed = torch.arange(1, H + 1, dtype=torch.float32).view(-1, 1).repeat(1, W) - x_embed = torch.arange(1, W + 1, dtype=torch.float32).view(1, -1).repeat(H, 1) - - # Normalize - eps = 1e-6 - y_embed = y_embed / (y_embed[-1:, :] + eps) * scale - x_embed = x_embed / (x_embed[:, -1:] + eps) * scale - - dim_t = torch.arange(half, dtype=torch.float32) - dim_t = temperature ** (2 * (dim_t // 2) / half) - - pos_x = x_embed[:, :, None] / dim_t # [H, W, half] - pos_y = y_embed[:, :, None] / dim_t - - pos_x = torch.stack((pos_x[:, :, 0::2].sin(), pos_x[:, :, 1::2].cos()), dim=3).flatten(2) - pos_y = torch.stack((pos_y[:, :, 0::2].sin(), pos_y[:, :, 1::2].cos()), dim=3).flatten(2) - - # [H, W, 256] → [1, 256, H, W] - pos = torch.cat((pos_y, pos_x), dim=2).permute(2, 0, 1).unsqueeze(0) - return pos - - -# ═══════════════════════════════════════════════════════════════════════════ -# Main -# ═══════════════════════════════════════════════════════════════════════════ - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--checkpoint", required=True) - parser.add_argument("--image", required=True) - parser.add_argument("--outdir", default="tests/ref_phase3") - parser.add_argument("--blocks", default="0,1,7,15,23,31", - help="Comma-separated block indices to dump detailed intermediates") - args = parser.parse_args() - - os.makedirs(args.outdir, exist_ok=True) - device = "cpu" - dump_blocks = set(int(b) for b in args.blocks.split(",")) - - # ── Load checkpoint ───────────────────────────────────────────────── - print("Loading checkpoint...") - ckpt = torch.load(args.checkpoint, map_location=device, weights_only=False) - if "model" in ckpt: - ckpt = ckpt["model"] - - vit_prefix = "detector.backbone.vision_backbone.trunk." - vit_sd = {k[len(vit_prefix):]: v for k, v in ckpt.items() if k.startswith(vit_prefix)} - print(f" Found {len(vit_sd)} ViT keys") - - # ── Step 3.4: RoPE ───────────────────────────────────────────────── - print("\n=== RoPE Frequencies ===") - - # Window RoPE (24x24, scale_pos=1.0) - window_freqs = compute_axial_cis(dim=64, end_x=24, end_y=24, theta=10000.0, scale_pos=1.0) - # Store as real: [576, 32, 2] (cos, sin) - window_real = torch.view_as_real(window_freqs) - save_tensor(os.path.join(args.outdir, "rope_window_real"), window_real) - - # Global RoPE (72x72, scale_pos=24/72=1/3) - global_freqs = compute_axial_cis(dim=64, end_x=72, end_y=72, theta=10000.0, scale_pos=24.0/72.0) - global_real = torch.view_as_real(global_freqs) - save_tensor(os.path.join(args.outdir, "rope_global_real"), global_real) - - # Also dump the checkpoint's freqs_cis for block 0 (window) and block 7 (global) - for bi in [0, 7]: - key = f"blocks.{bi}.attn.freqs_cis" - if key in vit_sd: - fc = vit_sd[key] - fc_real = torch.view_as_real(fc) if fc.is_complex() else fc - save_tensor(os.path.join(args.outdir, f"ckpt_freqs_cis_block{bi}"), fc_real) - - # ── Step 3.9: Sinusoidal PE ───────────────────────────────────────── - print("\n=== Sinusoidal PE ===") - for S, name in [(288, "pe_288"), (144, "pe_144"), (72, "pe_72"), (36, "pe_36")]: - pe = sinusoidal_pe_2d(S, S, num_pos_feats=256) - save_tensor(os.path.join(args.outdir, name), pe) - - # ── Load and preprocess image ───────────────────────────────────── - print(f"\n=== Preprocessing: {args.image} ===") - img = Image.open(args.image).convert("RGB") - transform = v2.Compose([ - v2.ToDtype(torch.uint8, scale=True), - v2.Resize(size=(1008, 1008)), - v2.ToDtype(torch.float32, scale=True), - v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), - ]) - img_tensor = v2.functional.to_image(img) - img_preprocessed = transform(img_tensor).unsqueeze(0).to(device) - save_tensor(os.path.join(args.outdir, "preprocessed"), img_preprocessed) - - # Also save the raw CHW float image before normalization (for debugging preprocessing) - raw_transform = v2.Compose([ - v2.ToDtype(torch.uint8, scale=True), - v2.Resize(size=(1008, 1008)), - v2.ToDtype(torch.float32, scale=True), - ]) - img_raw_float = raw_transform(v2.functional.to_image(img)).unsqueeze(0) - save_tensor(os.path.join(args.outdir, "preprocessed_raw_float"), img_raw_float) - - # ── Manual ViT forward pass ─────────────────────────────────────── - print("\n=== ViT Forward Pass ===") - with torch.no_grad(): - E = 1024 - NH = 16 - HD = 64 - depth = 32 - MLP_DIM = 4736 - WS = 24 - global_blocks = {7, 15, 23, 31} - - # Step 3.2: Patch embedding - patch_w = vit_sd["patch_embed.proj.weight"] # [1024, 3, 14, 14] - x = F.conv2d(img_preprocessed, patch_w, stride=14) # [1, 1024, 72, 72] - x = x.permute(0, 2, 3, 1) # [1, 72, 72, 1024] - save_tensor(os.path.join(args.outdir, "patch_embed"), x) - - # Step 3.3: Positional embedding (tiled) - pos_embed = vit_sd["pos_embed"] # [1, 577, 1024] - pos = get_abs_pos(pos_embed, has_cls_token=True, hw=(72, 72), tiling=True) - save_tensor(os.path.join(args.outdir, "pos_embed_tiled"), pos) - x = x + pos - save_tensor(os.path.join(args.outdir, "after_pos_embed"), x) - - # ln_pre - ln_pre_w = vit_sd["ln_pre.weight"] - ln_pre_b = vit_sd["ln_pre.bias"] - x = F.layer_norm(x, [E], ln_pre_w, ln_pre_b, eps=1e-5) - save_tensor(os.path.join(args.outdir, "after_ln_pre"), x) - - # Step 3.6 & 3.7: Transformer blocks - for blk_idx in range(depth): - prefix = f"blocks.{blk_idx}." - is_global = blk_idx in global_blocks - ws = 0 if is_global else WS - dump_this = blk_idx in dump_blocks - blk_prefix = f"block_{blk_idx}" - - # Pre-norm - n1_w = vit_sd[prefix + "norm1.weight"] - n1_b = vit_sd[prefix + "norm1.bias"] - shortcut = x - xn = F.layer_norm(x, [E], n1_w, n1_b, eps=1e-5) - - if dump_this: - save_tensor(os.path.join(args.outdir, f"{blk_prefix}_after_norm1"), xn) - - # Window partition - if ws > 0: - H, W = xn.shape[1], xn.shape[2] - xn, pad_hw = window_partition(xn, ws) - if dump_this: - save_tensor(os.path.join(args.outdir, f"{blk_prefix}_after_winpart"), xn) - - # Attention - B_cur, Hc, Wc, _ = xn.shape - L = Hc * Wc - qkv_w = vit_sd[prefix + "attn.qkv.weight"] - qkv_b = vit_sd[prefix + "attn.qkv.bias"] - qkv = F.linear(xn, qkv_w, qkv_b).reshape(B_cur, L, 3, NH, HD) - q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0) # [B, NH, L, HD] - - if dump_this: - save_tensor(os.path.join(args.outdir, f"{blk_prefix}_q_pre_rope"), q) - save_tensor(os.path.join(args.outdir, f"{blk_prefix}_k_pre_rope"), k) - - # RoPE - freqs_key = prefix + "attn.freqs_cis" - if freqs_key in vit_sd: - freqs = vit_sd[freqs_key] - q, k = apply_rotary_enc(q, k, freqs) - - if dump_this: - save_tensor(os.path.join(args.outdir, f"{blk_prefix}_q_post_rope"), q) - save_tensor(os.path.join(args.outdir, f"{blk_prefix}_k_post_rope"), k) - save_tensor(os.path.join(args.outdir, f"{blk_prefix}_v"), v) - - # Scaled dot-product attention - attn_out = F.scaled_dot_product_attention(q, k, v) - attn_out = attn_out.view(B_cur, NH, Hc, Wc, HD).permute(0, 2, 3, 1, 4).reshape(B_cur, Hc, Wc, E) - - if dump_this: - save_tensor(os.path.join(args.outdir, f"{blk_prefix}_attn_out"), attn_out) - - # Output projection - proj_w = vit_sd[prefix + "attn.proj.weight"] - proj_b = vit_sd[prefix + "attn.proj.bias"] - attn_out = F.linear(attn_out, proj_w, proj_b) - - if dump_this: - save_tensor(os.path.join(args.outdir, f"{blk_prefix}_proj_out"), attn_out) - - # Window unpartition - if ws > 0: - attn_out = window_unpartition(attn_out, ws, pad_hw, (H, W)) - - # Residual - x = shortcut + attn_out - - if dump_this: - save_tensor(os.path.join(args.outdir, f"{blk_prefix}_after_attn_residual"), x) - - # FFN - n2_w = vit_sd[prefix + "norm2.weight"] - n2_b = vit_sd[prefix + "norm2.bias"] - shortcut = x - xn = F.layer_norm(x, [E], n2_w, n2_b, eps=1e-5) - - fc1_w = vit_sd[prefix + "mlp.fc1.weight"] - fc1_b = vit_sd[prefix + "mlp.fc1.bias"] - fc2_w = vit_sd[prefix + "mlp.fc2.weight"] - fc2_b = vit_sd[prefix + "mlp.fc2.bias"] - h = F.linear(xn, fc1_w, fc1_b) - h = F.gelu(h) - h = F.linear(h, fc2_w, fc2_b) - - x = shortcut + h - - # Always save per-block output for selected blocks - save_tensor(os.path.join(args.outdir, f"block_{blk_idx}_out"), x) - - # Final output - save_tensor(os.path.join(args.outdir, "vit_final"), x) - feats = x.permute(0, 3, 1, 2) # [1, 1024, 72, 72] - save_tensor(os.path.join(args.outdir, "vit_output_bchw"), feats) - - # ── Neck (SimpleFPN) — detector path ────────────────────────────── - print("\n=== Neck (Detector) ===") - neck_prefix = "detector.backbone.vision_backbone.convs." - neck_sd = {k[len(neck_prefix):]: v for k, v in ckpt.items() if k.startswith(neck_prefix)} - - with torch.no_grad(): - # Scale 0 (4x) - s0 = F.conv_transpose2d(feats, neck_sd["0.dconv_2x2_0.weight"], neck_sd["0.dconv_2x2_0.bias"], stride=2) - save_tensor(os.path.join(args.outdir, "neck_det_s0_deconv1"), s0) - s0 = F.gelu(s0) - save_tensor(os.path.join(args.outdir, "neck_det_s0_gelu"), s0) - s0 = F.conv_transpose2d(s0, neck_sd["0.dconv_2x2_1.weight"], neck_sd["0.dconv_2x2_1.bias"], stride=2) - save_tensor(os.path.join(args.outdir, "neck_det_s0_deconv2"), s0) - s0 = F.conv2d(s0, neck_sd["0.conv_1x1.weight"], neck_sd["0.conv_1x1.bias"]) - save_tensor(os.path.join(args.outdir, "neck_det_s0_conv1x1"), s0) - s0 = F.conv2d(s0, neck_sd["0.conv_3x3.weight"], neck_sd["0.conv_3x3.bias"], padding=1) - save_tensor(os.path.join(args.outdir, "neck_det_0"), s0) - - # Scale 1 (2x) - s1 = F.conv_transpose2d(feats, neck_sd["1.dconv_2x2.weight"], neck_sd["1.dconv_2x2.bias"], stride=2) - save_tensor(os.path.join(args.outdir, "neck_det_s1_deconv"), s1) - s1 = F.conv2d(s1, neck_sd["1.conv_1x1.weight"], neck_sd["1.conv_1x1.bias"]) - save_tensor(os.path.join(args.outdir, "neck_det_s1_conv1x1"), s1) - s1 = F.conv2d(s1, neck_sd["1.conv_3x3.weight"], neck_sd["1.conv_3x3.bias"], padding=1) - save_tensor(os.path.join(args.outdir, "neck_det_1"), s1) - - # Scale 2 (1x) - s2 = F.conv2d(feats, neck_sd["2.conv_1x1.weight"], neck_sd["2.conv_1x1.bias"]) - save_tensor(os.path.join(args.outdir, "neck_det_s2_conv1x1"), s2) - s2 = F.conv2d(s2, neck_sd["2.conv_3x3.weight"], neck_sd["2.conv_3x3.bias"], padding=1) - save_tensor(os.path.join(args.outdir, "neck_det_2"), s2) - - # Scale 3 (0.5x) - s3 = F.max_pool2d(feats, kernel_size=2, stride=2) - save_tensor(os.path.join(args.outdir, "neck_det_s3_pool"), s3) - s3 = F.conv2d(s3, neck_sd["3.conv_1x1.weight"], neck_sd["3.conv_1x1.bias"]) - save_tensor(os.path.join(args.outdir, "neck_det_s3_conv1x1"), s3) - s3 = F.conv2d(s3, neck_sd["3.conv_3x3.weight"], neck_sd["3.conv_3x3.bias"], padding=1) - save_tensor(os.path.join(args.outdir, "neck_det_3"), s3) - - # ── Neck (SimpleFPN) — tracker path ────────────────────────────── - print("\n=== Neck (Tracker) ===") - trk_prefix = "detector.backbone.vision_backbone.sam2_convs." - trk_sd = {k[len(trk_prefix):]: v for k, v in ckpt.items() if k.startswith(trk_prefix)} - - if trk_sd: - with torch.no_grad(): - s0t = F.conv_transpose2d(feats, trk_sd["0.dconv_2x2_0.weight"], trk_sd["0.dconv_2x2_0.bias"], stride=2) - s0t = F.gelu(s0t) - s0t = F.conv_transpose2d(s0t, trk_sd["0.dconv_2x2_1.weight"], trk_sd["0.dconv_2x2_1.bias"], stride=2) - s0t = F.conv2d(s0t, trk_sd["0.conv_1x1.weight"], trk_sd["0.conv_1x1.bias"]) - s0t = F.conv2d(s0t, trk_sd["0.conv_3x3.weight"], trk_sd["0.conv_3x3.bias"], padding=1) - save_tensor(os.path.join(args.outdir, "neck_trk_0"), s0t) - - s1t = F.conv_transpose2d(feats, trk_sd["1.dconv_2x2.weight"], trk_sd["1.dconv_2x2.bias"], stride=2) - s1t = F.conv2d(s1t, trk_sd["1.conv_1x1.weight"], trk_sd["1.conv_1x1.bias"]) - s1t = F.conv2d(s1t, trk_sd["1.conv_3x3.weight"], trk_sd["1.conv_3x3.bias"], padding=1) - save_tensor(os.path.join(args.outdir, "neck_trk_1"), s1t) - - s2t = F.conv2d(feats, trk_sd["2.conv_1x1.weight"], trk_sd["2.conv_1x1.bias"]) - s2t = F.conv2d(s2t, trk_sd["2.conv_3x3.weight"], trk_sd["2.conv_3x3.bias"], padding=1) - save_tensor(os.path.join(args.outdir, "neck_trk_2"), s2t) - - s3t = F.max_pool2d(feats, kernel_size=2, stride=2) - s3t = F.conv2d(s3t, trk_sd["3.conv_1x1.weight"], trk_sd["3.conv_1x1.bias"]) - s3t = F.conv2d(s3t, trk_sd["3.conv_3x3.weight"], trk_sd["3.conv_3x3.bias"], padding=1) - save_tensor(os.path.join(args.outdir, "neck_trk_3"), s3t) - else: - print(" No tracker neck weights found (sam2_convs); skipping.") - - print("\n=== Done! ===") - print(f"All tensors saved to: {args.outdir}") - - -if __name__ == "__main__": - main() diff --git a/tests/dump_phase4_reference.py b/tests/dump_phase4_reference.py deleted file mode 100644 index f3efb39..0000000 --- a/tests/dump_phase4_reference.py +++ /dev/null @@ -1,257 +0,0 @@ -#!/usr/bin/env python3 -""" -Dump Phase 4 text encoder reference tensors for numerical comparison. - -Usage: - uv run python tests/dump_phase4_reference.py \ - --checkpoint raw_weights/sam3.pt \ - --tokenizer raw_weights/tokenizer.json \ - --prompts tests/phase4_prompts.tsv \ - --outdir tests/ref_phase4 -""" - -import argparse -import os -from typing import List, Tuple - -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F -from tokenizers import Tokenizer - - -def save_tensor(path: str, t: torch.Tensor) -> None: - t = t.detach().cpu().float().contiguous() - os.makedirs(os.path.dirname(path) or ".", exist_ok=True) - with open(path + ".bin", "wb") as f: - f.write(t.numpy().tobytes()) - with open(path + ".shape", "w") as f: - f.write(",".join(str(d) for d in t.shape)) - - -def save_i32_tensor(path: str, values: List[int]) -> None: - arr = np.asarray(values, dtype=np.int32) - os.makedirs(os.path.dirname(path) or ".", exist_ok=True) - with open(path + ".bin", "wb") as f: - f.write(arr.tobytes()) - with open(path + ".shape", "w") as f: - f.write(str(arr.shape[0])) - - -def load_prompt_cases(path: str) -> List[Tuple[str, str]]: - cases: List[Tuple[str, str]] = [] - with open(path, "r", encoding="utf-8") as f: - for line in f: - line = line.rstrip("\n") - if not line: - continue - if "\t" in line: - prompt_id, text = line.split("\t", 1) - else: - prompt_id, text = f"prompt_{len(cases):02d}", line - cases.append((prompt_id, text)) - return cases - - -def save_ggml_2d(path: str, t_td: torch.Tensor) -> None: - # ggml stores a [D, T] tensor with D contiguous. A contiguous PyTorch [T, D] - # tensor has the same byte order, so we write those bytes but record ggml dims. - assert t_td.ndim == 2 - t_td = t_td.detach().cpu().float().contiguous() - os.makedirs(os.path.dirname(path) or ".", exist_ok=True) - with open(path + ".bin", "wb") as f: - f.write(t_td.numpy().tobytes()) - with open(path + ".shape", "w") as f: - f.write(f"{t_td.shape[1]},{t_td.shape[0]}") - -def save_ggml_2d_batch_first(path: str, x: torch.Tensor) -> None: - assert x.ndim == 3 and x.shape[0] == 1 - save_ggml_2d(path, x[0]) - - -def save_ggml_2d_seq_first(path: str, x: torch.Tensor) -> None: - assert x.ndim == 3 and x.shape[1] == 1 - save_ggml_2d(path, x[:, 0, :]) - - -def tokenize_prompt(tokenizer: Tokenizer, text: str, context_length: int) -> List[int]: - sot_token = 49406 - eot_token = 49407 - - token_ids = list(tokenizer.encode(text).ids) - if not token_ids or token_ids[0] != sot_token: - token_ids = [sot_token] + token_ids - if token_ids[-1] != eot_token: - token_ids.append(eot_token) - - if len(token_ids) > context_length: - token_ids = token_ids[:context_length] - token_ids[-1] = eot_token - while len(token_ids) < context_length: - token_ids.append(0) - return token_ids - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--checkpoint", required=True) - parser.add_argument("--tokenizer", default="raw_weights/tokenizer.json") - parser.add_argument("--prompts", default="tests/phase4_prompts.tsv") - parser.add_argument("--outdir", default="tests/ref_phase4") - args = parser.parse_args() - - os.makedirs(args.outdir, exist_ok=True) - device = "cpu" - - print("Loading checkpoint...") - ckpt = torch.load(args.checkpoint, map_location=device, weights_only=False) - if "model" in ckpt: - ckpt = ckpt["model"] - - text_prefix = "detector.backbone.language_backbone.encoder." - text_sd = {k[len(text_prefix):]: v for k, v in ckpt.items() if k.startswith(text_prefix)} - resizer_prefix = "detector.backbone.language_backbone.resizer." - resizer_sd = { - k[len(resizer_prefix):]: v for k, v in ckpt.items() if k.startswith(resizer_prefix) - } - - tokenizer = Tokenizer.from_file(args.tokenizer) - prompts = load_prompt_cases(args.prompts) - - text_width = 1024 - text_heads = 16 - text_layers = 24 - context_length = 32 - - causal_mask = torch.full((context_length, context_length), float("-inf"), device=device) - causal_mask = torch.triu(causal_mask, diagonal=1) - - with torch.no_grad(): - for prompt_id, prompt_text in prompts: - prompt_dir = os.path.join(args.outdir, prompt_id) - os.makedirs(prompt_dir, exist_ok=True) - - token_ids = tokenize_prompt(tokenizer, prompt_text, context_length) - with open(os.path.join(prompt_dir, "prompt.txt"), "w", encoding="utf-8") as f: - f.write(prompt_text) - - save_i32_tensor(os.path.join(prompt_dir, "token_ids"), token_ids) - save_tensor(os.path.join(prompt_dir, "causal_mask"), causal_mask) - - tokens = torch.tensor([token_ids], dtype=torch.long, device=device) - - x_bf = F.embedding(tokens, text_sd["token_embedding.weight"]) - save_ggml_2d_batch_first(os.path.join(prompt_dir, "text_token_embed"), x_bf) - - x_bf = x_bf + text_sd["positional_embedding"][:context_length] - save_ggml_2d_batch_first(os.path.join(prompt_dir, "text_after_pos_embed"), x_bf) - - x = x_bf.transpose(0, 1).contiguous() # [T, 1, D] - - mha = nn.MultiheadAttention(text_width, text_heads, batch_first=False).to(device) - mha.eval() - - for i in range(text_layers): - prefix = f"transformer.resblocks.{i}." - - xn1 = F.layer_norm( - x, - (text_width,), - text_sd[prefix + "ln_1.weight"], - text_sd[prefix + "ln_1.bias"], - ) - save_ggml_2d( - os.path.join(prompt_dir, f"text_block_{i:02d}_after_ln1"), - xn1[:, 0, :], - ) - - qkv = F.linear( - xn1, - text_sd[prefix + "attn.in_proj_weight"], - text_sd[prefix + "attn.in_proj_bias"], - ) - save_ggml_2d( - os.path.join(prompt_dir, f"text_block_{i:02d}_qkv"), - qkv[:, 0, :], - ) - - mha.in_proj_weight.copy_(text_sd[prefix + "attn.in_proj_weight"]) - mha.in_proj_bias.copy_(text_sd[prefix + "attn.in_proj_bias"]) - mha.out_proj.weight.copy_(text_sd[prefix + "attn.out_proj.weight"]) - mha.out_proj.bias.copy_(text_sd[prefix + "attn.out_proj.bias"]) - - attn_out = mha(xn1, xn1, xn1, need_weights=False, attn_mask=causal_mask)[0] - save_ggml_2d( - os.path.join(prompt_dir, f"text_block_{i:02d}_attn_out"), - attn_out[:, 0, :], - ) - - x = x + attn_out - save_ggml_2d( - os.path.join(prompt_dir, f"text_block_{i:02d}_after_attn_residual"), - x[:, 0, :], - ) - - xn2 = F.layer_norm( - x, - (text_width,), - text_sd[prefix + "ln_2.weight"], - text_sd[prefix + "ln_2.bias"], - ) - save_ggml_2d( - os.path.join(prompt_dir, f"text_block_{i:02d}_after_ln2"), - xn2[:, 0, :], - ) - - mlp_fc1 = F.linear( - xn2, - text_sd[prefix + "mlp.c_fc.weight"], - text_sd[prefix + "mlp.c_fc.bias"], - ) - save_ggml_2d( - os.path.join(prompt_dir, f"text_block_{i:02d}_mlp_fc1"), - mlp_fc1[:, 0, :], - ) - - mlp_gelu = F.gelu(mlp_fc1) - save_ggml_2d( - os.path.join(prompt_dir, f"text_block_{i:02d}_mlp_gelu"), - mlp_gelu[:, 0, :], - ) - - mlp_out = F.linear( - mlp_gelu, - text_sd[prefix + "mlp.c_proj.weight"], - text_sd[prefix + "mlp.c_proj.bias"], - ) - save_ggml_2d( - os.path.join(prompt_dir, f"text_block_{i:02d}_mlp_out"), - mlp_out[:, 0, :], - ) - - x = x + mlp_out - save_ggml_2d( - os.path.join(prompt_dir, f"text_block_{i:02d}_out"), - x[:, 0, :], - ) - - x = F.layer_norm( - x, - (text_width,), - text_sd["ln_final.weight"], - text_sd["ln_final.bias"], - ) - save_ggml_2d_seq_first(os.path.join(prompt_dir, "text_final_ln"), x) - - text_features = F.linear(x, resizer_sd["weight"], resizer_sd["bias"]) - save_ggml_2d_seq_first( - os.path.join(prompt_dir, "text_features_2d"), - text_features, - ) - - print(f"[done] {prompt_id}: \"{prompt_text}\"") - - -if __name__ == "__main__": - main() diff --git a/tests/dump_phase5_reference.py b/tests/dump_phase5_reference.py deleted file mode 100644 index 5703671..0000000 --- a/tests/dump_phase5_reference.py +++ /dev/null @@ -1,778 +0,0 @@ -#!/usr/bin/env python3 -""" -Dump Phase 5 (Detector/PCS) reference tensors from SAM3 checkpoint for C++ verification. -Covers: Fusion Encoder, DETR Decoder, DotProductScoring, Segmentation Head. -Self-contained — builds forward pass from raw weights, no triton dependency. - -Usage: - uv run python tests/dump_phase5_reference.py --checkpoint raw_weights/sam3.pt --image tests/test_random.jpg --outdir tests/ref_phase5/ -""" -import argparse -import math -import os -import sys - -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F -from PIL import Image -from torchvision.transforms import v2 - - -def save_tensor(path, t): - t = t.detach().cpu().float().contiguous() - os.makedirs(os.path.dirname(path) or ".", exist_ok=True) - with open(path + ".bin", "wb") as f: - f.write(t.numpy().tobytes()) - with open(path + ".shape", "w") as f: - f.write(",".join(str(d) for d in t.shape)) - print(f" saved {path} shape={list(t.shape)}") - - -def save_tensor_i32(path, t): - t = t.detach().cpu().to(torch.int32).contiguous() - os.makedirs(os.path.dirname(path) or ".", exist_ok=True) - with open(path + ".bin", "wb") as f: - f.write(t.numpy().tobytes()) - with open(path + ".shape", "w") as f: - f.write(",".join(str(d) for d in t.shape)) - print(f" saved {path} shape={list(t.shape)} dtype=int32") - - -# ── Inlined helpers ───────────────────────────────────────────────────── - -def inverse_sigmoid(x, eps=1e-3): - x = x.clamp(min=0, max=1) - x1 = x.clamp(min=eps) - x2 = (1 - x).clamp(min=eps) - return torch.log(x1 / x2) - - -def gen_sineembed_for_position(pos_tensor, num_feats=256): - """Generate sinusoidal positional embedding for 2D/4D reference points.""" - assert num_feats % 2 == 0 - num_feats = num_feats // 2 - scale = 2 * math.pi - dim_t = torch.arange(num_feats, dtype=torch.float32, device=pos_tensor.device) - dim_t = 10000 ** (2 * (torch.div(dim_t, 2, rounding_mode="floor")) / num_feats) - x_embed = pos_tensor[:, :, 0] * scale - y_embed = pos_tensor[:, :, 1] * scale - pos_x = x_embed[:, :, None] / dim_t - pos_y = y_embed[:, :, None] / dim_t - pos_x = torch.stack((pos_x[:, :, 0::2].sin(), pos_x[:, :, 1::2].cos()), dim=3).flatten(2) - pos_y = torch.stack((pos_y[:, :, 0::2].sin(), pos_y[:, :, 1::2].cos()), dim=3).flatten(2) - if pos_tensor.size(-1) == 2: - pos = torch.cat((pos_y, pos_x), dim=2) - elif pos_tensor.size(-1) == 4: - w_embed = pos_tensor[:, :, 2] * scale - pos_w = w_embed[:, :, None] / dim_t - pos_w = torch.stack((pos_w[:, :, 0::2].sin(), pos_w[:, :, 1::2].cos()), dim=3).flatten(2) - h_embed = pos_tensor[:, :, 3] * scale - pos_h = h_embed[:, :, None] / dim_t - pos_h = torch.stack((pos_h[:, :, 0::2].sin(), pos_h[:, :, 1::2].cos()), dim=3).flatten(2) - pos = torch.cat((pos_y, pos_x, pos_w, pos_h), dim=2) - else: - raise ValueError(f"Unknown pos_tensor shape(-1):{pos_tensor.size(-1)}") - return pos - - -def box_cxcywh_to_xyxy(x): - x_c, y_c, w, h = x.unbind(-1) - b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c + 0.5 * w), (y_c + 0.5 * h)] - return torch.stack(b, dim=-1) - - -def mlp_forward(x, sd, prefix, num_layers, activation="relu"): - """Forward through MLP layers from state dict.""" - for i in range(num_layers): - w = sd[f"{prefix}.layers.{i}.weight"] - b = sd[f"{prefix}.layers.{i}.bias"] - x = F.linear(x, w, b) - if i < num_layers - 1: - x = F.relu(x) - return x - - -def multihead_attention_forward(query, key, value, w_q, b_q, w_k, b_k, w_v, b_v, - w_out, b_out, num_heads, attn_mask=None, - key_padding_mask=None): - """Manual multi-head attention forward.""" - D = query.shape[-1] - HD = D // num_heads - - Q = F.linear(query, w_q, b_q) - K = F.linear(key, w_k, b_k) - V = F.linear(value, w_v, b_v) - - # [BS, N, D] -> [BS, N, NH, HD] -> [BS, NH, N, HD] - BS = Q.shape[0] - N_q = Q.shape[1] - N_kv = K.shape[1] - - Q = Q.reshape(BS, N_q, num_heads, HD).permute(0, 2, 1, 3) - K = K.reshape(BS, N_kv, num_heads, HD).permute(0, 2, 1, 3) - V = V.reshape(BS, N_kv, num_heads, HD).permute(0, 2, 1, 3) - - combined_mask = attn_mask - if key_padding_mask is not None: - pad_bias = torch.zeros( - BS, 1, 1, N_kv, dtype=Q.dtype, device=Q.device - ) - pad_bias = pad_bias.masked_fill( - key_padding_mask[:, None, None, :], float("-inf") - ) - combined_mask = pad_bias if combined_mask is None else combined_mask + pad_bias - - # Attention - attn_output = F.scaled_dot_product_attention(Q, K, V, attn_mask=combined_mask) - - # [BS, NH, N_q, HD] -> [BS, N_q, D] - attn_output = attn_output.permute(0, 2, 1, 3).reshape(BS, N_q, D) - out = F.linear(attn_output, w_out, b_out) - return out - - -def fused_mha_forward(query, key, value, in_proj_w, in_proj_b, out_w, out_b, - num_heads, attn_mask=None, key_padding_mask=None): - """Multi-head attention with fused in_proj weights.""" - D = query.shape[-1] - w_q, w_k, w_v = in_proj_w[:D], in_proj_w[D:2*D], in_proj_w[2*D:] - b_q, b_k, b_v = in_proj_b[:D], in_proj_b[D:2*D], in_proj_b[2*D:] - return multihead_attention_forward(query, key, value, w_q, b_q, w_k, b_k, w_v, b_v, - out_w, out_b, num_heads, attn_mask, key_padding_mask) - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--checkpoint", required=True) - parser.add_argument("--image", required=True) - parser.add_argument("--outdir", default="tests/ref_phase5") - parser.add_argument("--text", default="yellow school bus") - args = parser.parse_args() - - os.makedirs(args.outdir, exist_ok=True) - device = "cpu" - D = 256 - NH = 8 - HD = D // NH - NQ = 200 - T = 32 # text context length - H = 72 # image feature spatial dim (1008 / 14) - - print("Loading checkpoint...") - ckpt = torch.load(args.checkpoint, map_location=device, weights_only=False) - if "model" in ckpt: - ckpt = ckpt["model"] - - # ── Load required weights from specific modules ───────────────────── - - # We need the image features & text features as inputs. - # To get them, we first need to run the ViT backbone + neck + text encoder. - # These were already verified in prior phases, so we load cached reference tensors. - # But we need fresh tensors to match exact C++ inputs. - - # Let's first generate fresh image features using dump_vit_reference approach - print(f"Loading image: {args.image}") - img = Image.open(args.image).convert("RGB") - transform = v2.Compose([ - v2.ToDtype(torch.uint8, scale=True), - v2.Resize(size=(1008, 1008)), - v2.ToDtype(torch.float32, scale=True), - v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), - ]) - img_tensor = v2.functional.to_image(img) - img_preprocessed = transform(img_tensor).unsqueeze(0).to(device) - - # ── ViT forward (reuse from dump_vit_reference.py) ───────────────── - # ... (this is too long to inline, we'll load pre-computed neck features) - # Instead, we'll load from the reference directory if available - ref_dir = "tests/ref" - neck_det_path = os.path.join(ref_dir, "neck_det_2") - if os.path.exists(neck_det_path + ".bin"): - print("Loading pre-computed neck features from tests/ref/") - # Load neck features - def load_tensor(path): - shape = list(map(int, open(path + ".shape").read().strip().split(","))) - data = np.fromfile(path + ".bin", dtype=np.float32).reshape(shape) - return torch.tensor(data) - - # Neck det features: [1, 256, H, H] for scale 2 (72x72) - neck_det_2 = load_tensor(os.path.join(ref_dir, "neck_det_2")) # [1, 256, 72, 72] - neck_det_1 = load_tensor(os.path.join(ref_dir, "neck_det_1")) # [1, 256, 144, 144] - neck_det_0 = load_tensor(os.path.join(ref_dir, "neck_det_0")) # [1, 256, 288, 288] - else: - print("ERROR: Pre-computed neck features not found. Run dump_vit_reference.py first.") - sys.exit(1) - - # ── Text encoder forward ──────────────────────────────────────────── - # For simplicity, we need text features. The text encoder was verified in phase 4. - # We can compute them here with the same approach. - # For now, load from existing ref or compute inline. - - # Actually, let's compute text features inline since we need them. - # The text encoder is: embedding -> transformer -> projection - - # Tokenize - from tokenizers import Tokenizer - tokenizer_path = os.path.join(os.path.dirname(args.checkpoint), "tokenizer.json") - if os.path.exists(tokenizer_path): - tokenizer = Tokenizer.from_file(tokenizer_path) - else: - tokenizer = Tokenizer.from_file("raw_weights/tokenizer.json") - - encoded = tokenizer.encode(args.text) - # tokenizer already includes SOT/EOT - token_ids = list(encoded.ids) - while len(token_ids) < T: - token_ids.append(0) # pad - token_ids = token_ids[:T] - - tokens_tensor = torch.tensor([token_ids], dtype=torch.long, device=device) # [1, T] - print(f" Token IDs: {token_ids[:10]}...") - save_tensor_i32(os.path.join(args.outdir, "token_ids"), tokens_tensor) - with open(os.path.join(args.outdir, "prompt.txt"), "w") as f: - f.write(args.text + "\n") - - # Text encoder forward - text_prefix = "detector.backbone.language_backbone.encoder." - text_sd = {k[len(text_prefix):]: v for k, v in ckpt.items() if k.startswith(text_prefix)} - text_width = 1024 - text_heads = 16 - text_layers = 24 - - with torch.no_grad(): - # Token + positional embedding - tok_emb = text_sd["token_embedding.weight"] # [49408, 1024] - pos_emb = text_sd["positional_embedding"] # [32, 1024] - x = tok_emb[tokens_tensor] + pos_emb[:T] # [1, T, 1024] - x = x.permute(1, 0, 2) # [T, 1, 1024] - - # Causal mask - causal_mask = torch.full((T, T), float("-inf"), device=device) - causal_mask = torch.triu(causal_mask, diagonal=1) - - # Transformer blocks - for i in range(text_layers): - p = f"transformer.resblocks.{i}." - # Pre-norm SA - xn = F.layer_norm(x, [text_width], text_sd[p + "ln_1.weight"], text_sd[p + "ln_1.bias"]) - # Fused in_proj - qkv_w = text_sd[p + "attn.in_proj_weight"] - qkv_b = text_sd[p + "attn.in_proj_bias"] - out_w = text_sd[p + "attn.out_proj.weight"] - out_b = text_sd[p + "attn.out_proj.bias"] - - # Manual MHA (T, BS, D) -> seq-first - xn_bf = xn.permute(1, 0, 2) # [1, T, 1024] - sa_out = fused_mha_forward(xn_bf, xn_bf, xn_bf, qkv_w, qkv_b, out_w, out_b, - text_heads, attn_mask=causal_mask) - x = x + sa_out.permute(1, 0, 2) - - # Pre-norm FFN - xn = F.layer_norm(x, [text_width], text_sd[p + "ln_2.weight"], text_sd[p + "ln_2.bias"]) - fc1_w = text_sd[p + "mlp.c_fc.weight"] - fc1_b = text_sd[p + "mlp.c_fc.bias"] - fc2_w = text_sd[p + "mlp.c_proj.weight"] - fc2_b = text_sd[p + "mlp.c_proj.bias"] - h = F.linear(xn, fc1_w, fc1_b) - h = F.gelu(h, approximate="tanh") - h = F.linear(h, fc2_w, fc2_b) - x = x + h - - # Final LN - x = F.layer_norm(x, [text_width], text_sd["ln_final.weight"], text_sd["ln_final.bias"]) - # x: [T, 1, 1024] - - # Project to 256-dim via resizer (Linear(1024→256)) - resizer_prefix = "detector.backbone.language_backbone.resizer." - resizer_sd = {k[len(resizer_prefix):]: v for k, v in ckpt.items() if k.startswith(resizer_prefix)} - txt_projected = F.linear(x, resizer_sd["weight"], resizer_sd.get("bias")) - # txt_projected: [T, 1, 256] - - save_tensor(os.path.join(args.outdir, "text_features"), txt_projected) - print(f" Text features: {list(txt_projected.shape)}") - - # Positional encoding for image features (sinusoidal) - # Compute sinusoidal PE manually - def compute_sine_pe(h, w, num_pos_feats=256, temperature=10000, normalize=True, scale=2*math.pi): - """Sinusoidal PE matching SAM3's PositionEmbeddingSine.""" - not_mask = torch.ones(1, h, w, dtype=torch.float32) - y_embed = not_mask.cumsum(1, dtype=torch.float32) - x_embed = not_mask.cumsum(2, dtype=torch.float32) - if normalize: - eps = 1e-6 - y_embed = (y_embed - 0.5) / (y_embed[:, -1:, :] + eps) * scale - x_embed = (x_embed - 0.5) / (x_embed[:, :, -1:] + eps) * scale - - dim_t = torch.arange(num_pos_feats, dtype=torch.float32) - dim_t = temperature ** (2 * (dim_t // 2) / num_pos_feats) - - pos_x = x_embed[:, :, :, None] / dim_t # [1, H, W, D/2] - pos_y = y_embed[:, :, :, None] / dim_t - pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3) - pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3) - pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) # [1, D, H, W] - return pos - - img_pe_72 = compute_sine_pe(H, H, num_pos_feats=D//2) # [1, 256, 72, 72] - save_tensor(os.path.join(args.outdir, "img_pe_72"), img_pe_72) - - # ═══════════════════════════════════════════════════════════════════════ - # Step 5.1: Fusion Encoder (6 layers) - # ═══════════════════════════════════════════════════════════════════════ - print("\n=== Fusion Encoder ===") - - fenc_prefix = "detector.transformer.encoder." - fenc_sd = {k[len("detector.transformer."):]: v for k, v in ckpt.items() - if k.startswith("detector.transformer.encoder.")} - - with torch.no_grad(): - # Image features: [1, 256, 72, 72] -> flatten to [1, 5184, 256] (batch-first) - img_feat_flat = neck_det_2.flatten(2).permute(0, 2, 1) # [1, 5184, 256] - img_pe_flat = img_pe_72.flatten(2).permute(0, 2, 1) # [1, 5184, 256] - - # Text features: [T, 1, 256] -> [1, T, 256] for batch-first - prompt_bf = txt_projected.permute(1, 0, 2) # [1, T, 256] - - # Text mask: True for padding (0 tokens), False for valid - txt_mask = torch.tensor([[tid == 0 for tid in token_ids]], dtype=torch.bool) # [1, T] - valid_scale = float(T) / max(1, sum(tid != 0 for tid in token_ids)) - text_valid_mask = torch.tensor( - [[[valid_scale if tid != 0 else 0.0] for tid in token_ids]], - dtype=torch.float32, - ).permute(1, 2, 0) # [T, 1, 1] - - save_tensor(os.path.join(args.outdir, "fenc_img_input"), img_feat_flat) - save_tensor(os.path.join(args.outdir, "fenc_pos_embed"), img_pe_flat) - save_tensor(os.path.join(args.outdir, "fenc_prompt"), prompt_bf) - save_tensor(os.path.join(args.outdir, "text_valid_mask"), text_valid_mask) - - # Forward through 6 fusion encoder layers - output = img_feat_flat # [1, HW, D] batch-first - - for layer_idx in range(6): - lp = f"encoder.layers.{layer_idx}." - - # 1. Self-attention (pre-norm, pos_enc_at_attn=True) - shortcut = output - xn = F.layer_norm(output, [D], - fenc_sd[lp + "norm1.weight"], fenc_sd[lp + "norm1.bias"]) - q = k = xn + img_pe_flat # add pos to Q and K - - # SA weights - sa_qkv_w = fenc_sd[lp + "self_attn.in_proj_weight"] - sa_qkv_b = fenc_sd[lp + "self_attn.in_proj_bias"] - sa_out_w = fenc_sd[lp + "self_attn.out_proj.weight"] - sa_out_b = fenc_sd[lp + "self_attn.out_proj.bias"] - - sa_out = fused_mha_forward(q, k, xn, sa_qkv_w, sa_qkv_b, sa_out_w, sa_out_b, NH) - output = shortcut + sa_out - - # 2. Cross-attention (pre-norm, no pos at CA) - shortcut = output - xn = F.layer_norm(output, [D], - fenc_sd[lp + "norm2.weight"], fenc_sd[lp + "norm2.bias"]) - # cross_attn_image: Q from image, K/V from text - ca_qkv_w = fenc_sd[lp + "cross_attn_image.in_proj_weight"] - ca_qkv_b = fenc_sd[lp + "cross_attn_image.in_proj_bias"] - ca_out_w = fenc_sd[lp + "cross_attn_image.out_proj.weight"] - ca_out_b = fenc_sd[lp + "cross_attn_image.out_proj.bias"] - - # Split in_proj into separate Q/K/V for cross-attention - w_q_ca = ca_qkv_w[:D] - w_k_ca = ca_qkv_w[D:2*D] - w_v_ca = ca_qkv_w[2*D:] - b_q_ca = ca_qkv_b[:D] - b_k_ca = ca_qkv_b[D:2*D] - b_v_ca = ca_qkv_b[2*D:] - - ca_out = multihead_attention_forward( - xn, prompt_bf, prompt_bf, - w_q_ca, b_q_ca, w_k_ca, b_k_ca, w_v_ca, b_v_ca, - ca_out_w, ca_out_b, NH, - key_padding_mask=txt_mask) - - output = shortcut + ca_out - - # 3. FFN (pre-norm, relu activation) - shortcut = output - xn = F.layer_norm(output, [D], - fenc_sd[lp + "norm3.weight"], fenc_sd[lp + "norm3.bias"]) - h = F.linear(xn, fenc_sd[lp + "linear1.weight"], fenc_sd[lp + "linear1.bias"]) - h = F.relu(h) - h = F.linear(h, fenc_sd[lp + "linear2.weight"], fenc_sd[lp + "linear2.bias"]) - output = shortcut + h - - save_tensor(os.path.join(args.outdir, f"fenc_layer{layer_idx}_out"), output) - - # Transpose to seq-first for downstream: [HW, 1, D] - fenc_output_sf = output.permute(1, 0, 2) # [5184, 1, D] seq-first - save_tensor(os.path.join(args.outdir, "fenc_output"), fenc_output_sf) - save_tensor(os.path.join(args.outdir, "fenc_output_bf"), output) # batch-first for decoder - print(f" fenc_output shape: {list(fenc_output_sf.shape)}") - - # ═══════════════════════════════════════════════════════════════════════ - # Step 5.2: DETR Decoder (6 layers) - # ═══════════════════════════════════════════════════════════════════════ - print("\n=== DETR Decoder ===") - - ddec_prefix = "detector.transformer.decoder." - ddec_sd = {k[len("detector.transformer."):]: v for k, v in ckpt.items() - if k.startswith("detector.transformer.decoder.")} - - with torch.no_grad(): - # Query embeddings and reference points - query_embed = ddec_sd["decoder.query_embed.weight"] # [NQ, D] - ref_pts_raw = ddec_sd["decoder.reference_points.weight"] # [NQ, 4] - reference_boxes = ref_pts_raw.sigmoid().unsqueeze(1) # [NQ, 1, 4] - - save_tensor(os.path.join(args.outdir, "ddec_query_embed"), query_embed) - save_tensor(os.path.join(args.outdir, "ddec_ref_pts_raw"), ref_pts_raw) - save_tensor(os.path.join(args.outdir, "ddec_ref_boxes_init"), reference_boxes) - - # Presence token - presence_token_w = ddec_sd["decoder.presence_token.weight"] # [1, D] - presence_out = presence_token_w[None].expand(1, 1, -1) # [1, 1, D] - save_tensor(os.path.join(args.outdir, "ddec_presence_token"), presence_out) - - # Encoder output is [HW, BS, D] (seq-first) for the decoder - enc_memory = fenc_output_sf # [5184, 1, D] - enc_pos = img_pe_flat.permute(1, 0, 2) # [5184, 1, D] seq-first - - # valid_ratios: [1, 1, 2] = ones (no masking) - valid_ratios = torch.ones(1, 1, 2, device=device) - spatial_shapes = torch.tensor([[H, H]], dtype=torch.long) - - # Init output (tgt) - output = query_embed.unsqueeze(1).repeat(1, 1, 1) # [NQ, 1, D] - - # Text features for cross-attention: seq-first [T, 1, D] - prompt_sf = txt_projected # already [T, 1, D] - - # ── RPB coordinates ───────────────────────────────────────── - coords_h = torch.arange(0, H, dtype=torch.float32) / H - coords_w = torch.arange(0, H, dtype=torch.float32) / H - - for layer_idx in range(6): - lp = f"decoder.layers.{layer_idx}." - - # Reference points input - reference_points_input = reference_boxes[:, :, None] * torch.cat([valid_ratios, valid_ratios], -1)[None, :] - query_sine_embed = gen_sineembed_for_position(reference_points_input[:, :, 0, :], D) - # query_sine_embed: [NQ, 1, 512] - - # Conditional query pos: ref_point_head MLP (512→256→256) - query_pos = mlp_forward(query_sine_embed, ddec_sd, - "decoder.ref_point_head", 2) # [NQ, 1, D] - - if layer_idx == 0: - save_tensor(os.path.join(args.outdir, "ddec_query_pos_0"), query_pos) - save_tensor(os.path.join(args.outdir, "ddec_query_sine_0"), query_sine_embed) - - # ── Box RPB ───────────────────────────────────────────── - boxes_xyxy = box_cxcywh_to_xyxy(reference_boxes).permute(1, 0, 2) # [1, NQ, 4] - bs_d, nq_d, _ = boxes_xyxy.shape - - deltas_y = coords_h.view(1, -1, 1) - boxes_xyxy.reshape(-1, 1, 4)[:, :, 1:4:2] - deltas_y = deltas_y.view(bs_d, nq_d, -1, 2) # [1, NQ, H, 2] - deltas_x = coords_w.view(1, -1, 1) - boxes_xyxy.reshape(-1, 1, 4)[:, :, 0:3:2] - deltas_x = deltas_x.view(bs_d, nq_d, -1, 2) # [1, NQ, W, 2] - - # Log transform - deltas_x_log = deltas_x * 8 - deltas_x_log = torch.sign(deltas_x_log) * torch.log2(torch.abs(deltas_x_log) + 1.0) / np.log2(8) - deltas_y_log = deltas_y * 8 - deltas_y_log = torch.sign(deltas_y_log) * torch.log2(torch.abs(deltas_y_log) + 1.0) / np.log2(8) - deltas_x = deltas_x_log - deltas_y = deltas_y_log - - # RPB MLP: [1, NQ, W, 2] → [1, NQ, W, NH] - rpb_x = mlp_forward(deltas_x, ddec_sd, "decoder.boxRPB_embed_x", 2) - rpb_y = mlp_forward(deltas_y, ddec_sd, "decoder.boxRPB_embed_y", 2) - - # Outer sum: [1, NQ, H, W, NH] - B_rpb = rpb_y.unsqueeze(3) + rpb_x.unsqueeze(2) - B_rpb = B_rpb.flatten(2, 3) # [1, NQ, H*W, NH] - B_rpb = B_rpb.permute(0, 3, 1, 2).contiguous() # [1, NH, NQ, H*W] - - # For presence token: prepend zeros - pres_rpb = torch.zeros(1, NH, 1, H*H, device=device) - memory_mask = torch.cat([pres_rpb, B_rpb], dim=2) # [1, NH, NQ+1, H*W] - memory_mask = memory_mask.flatten(0, 1) # [NH, NQ+1, H*W] - - if layer_idx == 0: - save_tensor(os.path.join(args.outdir, "ddec_rpb_mask_0"), memory_mask) - - # ── Self-attention ────────────────────────────────────── - # Prepend presence token - tgt_with_pres = torch.cat([presence_out, output], dim=0) # [NQ+1, 1, D] - qpos_with_pres = torch.cat([torch.zeros_like(presence_out), query_pos], dim=0) - - q_sa = k_sa = tgt_with_pres + qpos_with_pres - - sa_w = ddec_sd[lp + "self_attn.in_proj_weight"] - sa_b = ddec_sd[lp + "self_attn.in_proj_bias"] - sa_ow = ddec_sd[lp + "self_attn.out_proj.weight"] - sa_ob = ddec_sd[lp + "self_attn.out_proj.bias"] - - # Convert to batch-first for MHA - q_sa_bf = q_sa.permute(1, 0, 2) # [1, NQ+1, D] - k_sa_bf = k_sa.permute(1, 0, 2) - v_sa_bf = tgt_with_pres.permute(1, 0, 2) - - sa_out = fused_mha_forward(q_sa_bf, k_sa_bf, v_sa_bf, sa_w, sa_b, sa_ow, sa_ob, NH) - tgt_with_pres = tgt_with_pres + sa_out.permute(1, 0, 2) - tgt_with_pres = F.layer_norm(tgt_with_pres, [D], - ddec_sd[lp + "norm2.weight"], ddec_sd[lp + "norm2.bias"]) - if layer_idx == 0: - save_tensor(os.path.join(args.outdir, "ddec_layer0_after_sa"), tgt_with_pres) - - # ── Text cross-attention ───────────────────────────────── - tgt_q_ca_text = tgt_with_pres + qpos_with_pres - ca_text_w = ddec_sd[lp + "ca_text.in_proj_weight"] - ca_text_b = ddec_sd[lp + "ca_text.in_proj_bias"] - ca_text_ow = ddec_sd[lp + "ca_text.out_proj.weight"] - ca_text_ob = ddec_sd[lp + "ca_text.out_proj.bias"] - - # Q from queries, K/V from text - w_q_ct = ca_text_w[:D] - w_k_ct = ca_text_w[D:2*D] - w_v_ct = ca_text_w[2*D:] - b_q_ct = ca_text_b[:D] - b_k_ct = ca_text_b[D:2*D] - b_v_ct = ca_text_b[2*D:] - - tq_bf = tgt_q_ca_text.permute(1, 0, 2) # [1, NQ+1, D] - p_bf = prompt_sf.permute(1, 0, 2) # [1, T, D] - ct_out = multihead_attention_forward(tq_bf, p_bf, p_bf, - w_q_ct, b_q_ct, w_k_ct, b_k_ct, w_v_ct, b_v_ct, - ca_text_ow, ca_text_ob, NH, - key_padding_mask=txt_mask) - tgt_with_pres = tgt_with_pres + ct_out.permute(1, 0, 2) - tgt_with_pres = F.layer_norm(tgt_with_pres, [D], - ddec_sd[lp + "catext_norm.weight"], - ddec_sd[lp + "catext_norm.bias"]) - if layer_idx == 0: - save_tensor(os.path.join(args.outdir, "ddec_layer0_after_text_ca"), tgt_with_pres) - - # ── Image cross-attention with RPB ────────────────────── - tgt_q_ca_img = tgt_with_pres + qpos_with_pres - mem_k = enc_memory + enc_pos # [5184, 1, D] - - ca_w = ddec_sd[lp + "cross_attn.in_proj_weight"] - ca_b = ddec_sd[lp + "cross_attn.in_proj_bias"] - ca_ow = ddec_sd[lp + "cross_attn.out_proj.weight"] - ca_ob = ddec_sd[lp + "cross_attn.out_proj.bias"] - - w_q_ci = ca_w[:D] - w_k_ci = ca_w[D:2*D] - w_v_ci = ca_w[2*D:] - b_q_ci = ca_b[:D] - b_k_ci = ca_b[D:2*D] - b_v_ci = ca_b[2*D:] - - # Batch first - tqi_bf = tgt_q_ca_img.permute(1, 0, 2) # [1, NQ+1, D] - mk_bf = mem_k.permute(1, 0, 2) # [1, 5184, D] - mv_bf = enc_memory.permute(1, 0, 2) # [1, 5184, D] - - # Need to handle RPB mask as attn_mask - # memory_mask: [NH, NQ+1, H*W] — for PyTorch SDPA, need [BS*NH, NQ+1, H*W] - # For BS=1, just use as is - ci_out = multihead_attention_forward(tqi_bf, mk_bf, mv_bf, - w_q_ci, b_q_ci, w_k_ci, b_k_ci, w_v_ci, b_v_ci, - ca_ow, ca_ob, NH, - attn_mask=memory_mask) - tgt_with_pres = tgt_with_pres + ci_out.permute(1, 0, 2) - tgt_with_pres = F.layer_norm(tgt_with_pres, [D], - ddec_sd[lp + "norm1.weight"], ddec_sd[lp + "norm1.bias"]) - if layer_idx == 0: - save_tensor(os.path.join(args.outdir, "ddec_layer0_after_img_ca"), tgt_with_pres) - - # ── FFN ───────────────────────────────────────────────── - ffn = F.linear(tgt_with_pres, ddec_sd[lp + "linear1.weight"], ddec_sd[lp + "linear1.bias"]) - ffn = F.relu(ffn) - ffn = F.linear(ffn, ddec_sd[lp + "linear2.weight"], ddec_sd[lp + "linear2.bias"]) - tgt_with_pres = tgt_with_pres + ffn - tgt_with_pres = F.layer_norm(tgt_with_pres, [D], - ddec_sd[lp + "norm3.weight"], ddec_sd[lp + "norm3.bias"]) - if layer_idx == 0: - save_tensor(os.path.join(args.outdir, "ddec_layer0_full_out"), tgt_with_pres) - - # Split presence and queries - presence_out = tgt_with_pres[:1] # [1, 1, D] - output = tgt_with_pres[1:] # [NQ, 1, D] - - # Box refinement - reference_before_sigmoid = inverse_sigmoid(reference_boxes) - normed_output = F.layer_norm(output, [D], - ddec_sd["decoder.norm.weight"], ddec_sd["decoder.norm.bias"]) - delta_unsig = mlp_forward(normed_output, ddec_sd, "decoder.bbox_embed", 3) - outputs_unsig = delta_unsig + reference_before_sigmoid - new_reference_points = outputs_unsig.sigmoid() - reference_boxes = new_reference_points.detach() - - save_tensor(os.path.join(args.outdir, f"ddec_layer{layer_idx}_out"), output) - save_tensor(os.path.join(args.outdir, f"ddec_layer{layer_idx}_refboxes"), reference_boxes) - if layer_idx == 0: - save_tensor(os.path.join(args.outdir, f"ddec_layer{layer_idx}_presence"), presence_out) - - # Final normalization - normed_output = F.layer_norm(output, [D], - ddec_sd["decoder.norm.weight"], ddec_sd["decoder.norm.bias"]) - save_tensor(os.path.join(args.outdir, "ddec_normed_output"), normed_output) - save_tensor(os.path.join(args.outdir, "ddec_pred_boxes"), reference_boxes) - - # Presence head - pres_normed = F.layer_norm(presence_out, [D], - ddec_sd["decoder.presence_token_out_norm.weight"], - ddec_sd["decoder.presence_token_out_norm.bias"]) - pres_logit = mlp_forward(pres_normed, ddec_sd, "decoder.presence_token_head", 3) - save_tensor(os.path.join(args.outdir, "ddec_presence_logit"), pres_logit) - print(f" Presence logit: {pres_logit.item():.4f}") - - # ═══════════════════════════════════════════════════════════════════════ - # Step 5.2b: DotProductScoring - # ═══════════════════════════════════════════════════════════════════════ - print("\n=== DotProductScoring ===") - - scoring_prefix = "detector.dot_prod_scoring." - scoring_sd = {k[len(scoring_prefix):]: v for k, v in ckpt.items() - if k.startswith(scoring_prefix)} - - with torch.no_grad(): - # prompt_mlp: residual MLP + LayerNorm - prompt_for_scoring = prompt_sf.clone() # [T, 1, D] - prompt_mlp_out = prompt_for_scoring - orig = prompt_mlp_out.clone() - prompt_mlp_out = F.linear(prompt_mlp_out, scoring_sd["prompt_mlp.layers.0.weight"], - scoring_sd["prompt_mlp.layers.0.bias"]) - prompt_mlp_out = F.relu(prompt_mlp_out) - prompt_mlp_out = F.linear(prompt_mlp_out, scoring_sd["prompt_mlp.layers.1.weight"], - scoring_sd["prompt_mlp.layers.1.bias"]) - prompt_mlp_out = prompt_mlp_out + orig # residual - prompt_mlp_out = F.layer_norm(prompt_mlp_out, [D], - scoring_sd["prompt_mlp.out_norm.weight"], - scoring_sd["prompt_mlp.out_norm.bias"]) - save_tensor(os.path.join(args.outdir, "scoring_prompt_mlp_out"), prompt_mlp_out) - - # Mean pool text - # prompt_mask: True for padding, False for valid - is_valid = (~txt_mask).float().permute(1, 0)[..., None] # [T, 1, 1] - num_valid = torch.clamp(is_valid.sum(dim=0), min=1.0) # [1, 1] - pooled = (prompt_mlp_out * is_valid).sum(dim=0) / num_valid # [1, D] - save_tensor(os.path.join(args.outdir, "scoring_pooled"), pooled) - - # Project - proj_pooled = F.linear(pooled, scoring_sd["prompt_proj.weight"], - scoring_sd["prompt_proj.bias"]) - save_tensor(os.path.join(args.outdir, "scoring_proj_pooled"), proj_pooled) - - # Project hs - hs = normed_output.permute(1, 0, 2).unsqueeze(0) # [1, 1, NQ, D] - proj_hs = F.linear(hs, scoring_sd["hs_proj.weight"], scoring_sd["hs_proj.bias"]) - save_tensor(os.path.join(args.outdir, "scoring_proj_hs"), proj_hs) - - # Dot product scores - scores = torch.matmul(proj_hs, proj_pooled.unsqueeze(-1)) # [1, 1, NQ, 1] - scores = scores * (1.0 / math.sqrt(D)) - scores = scores.clamp(-12.0, 12.0) - save_tensor(os.path.join(args.outdir, "scoring_class_scores"), scores) - print(f" scores shape: {list(scores.shape)}, range: [{scores.min():.4f}, {scores.max():.4f}]") - - # ═══════════════════════════════════════════════════════════════════════ - # Step 5.3: Segmentation Head - # ═══════════════════════════════════════════════════════════════════════ - print("\n=== Segmentation Head ===") - - seg_prefix = "detector.segmentation_head." - seg_sd = {k[len(seg_prefix):]: v for k, v in ckpt.items() - if k.startswith(seg_prefix)} - - with torch.no_grad(): - # Cross-attend encoder hidden states to prompt - enc_hs = fenc_output_sf.clone() # [5184, 1, D] - tgt2 = F.layer_norm(enc_hs, [D], seg_sd["cross_attn_norm.weight"], seg_sd["cross_attn_norm.bias"]) - - ca_w = seg_sd["cross_attend_prompt.in_proj_weight"] - ca_b = seg_sd["cross_attend_prompt.in_proj_bias"] - ca_ow = seg_sd["cross_attend_prompt.out_proj.weight"] - ca_ob = seg_sd["cross_attend_prompt.out_proj.bias"] - - w_q = ca_w[:D]; w_k = ca_w[D:2*D]; w_v = ca_w[2*D:] - b_q = ca_b[:D]; b_k = ca_b[D:2*D]; b_v = ca_b[2*D:] - - tgt2_bf = tgt2.permute(1, 0, 2) - p_bf = prompt_sf.permute(1, 0, 2) - ca_out = multihead_attention_forward(tgt2_bf, p_bf, p_bf, - w_q, b_q, w_k, b_k, w_v, b_v, - ca_ow, ca_ob, NH, - key_padding_mask=txt_mask) - enc_hs = enc_hs + ca_out.permute(1, 0, 2) - save_tensor(os.path.join(args.outdir, "seg_enc_after_ca"), enc_hs) - - # Replace lowest-res FPN feat with encoder visual output - enc_visual = enc_hs.permute(1, 2, 0).reshape(1, D, H, H) - save_tensor(os.path.join(args.outdir, "seg_enc_visual"), enc_visual) - - # Pixel decoder - modified_feats = [neck_det_0.clone(), neck_det_1.clone(), enc_visual] - - prev_fpn = modified_feats[-1] # [1, D, 72, 72] - for pd_idx, bb_feat in enumerate(modified_feats[:-1][::-1]): - curr_fpn = bb_feat - prev_fpn = curr_fpn + F.interpolate(prev_fpn, size=curr_fpn.shape[-2:], mode="nearest") - prev_fpn = F.conv2d(prev_fpn, seg_sd[f"pixel_decoder.conv_layers.{pd_idx}.weight"], - seg_sd[f"pixel_decoder.conv_layers.{pd_idx}.bias"], padding=1) - prev_fpn = F.group_norm(prev_fpn, 8, - seg_sd[f"pixel_decoder.norms.{pd_idx}.weight"], - seg_sd[f"pixel_decoder.norms.{pd_idx}.bias"]) - prev_fpn = F.relu(prev_fpn) - save_tensor(os.path.join(args.outdir, f"seg_pixel_dec_stage{pd_idx}"), prev_fpn) - - pixel_embed = prev_fpn # [1, D, 288, 288] - save_tensor(os.path.join(args.outdir, "seg_pixel_decoder_out"), pixel_embed) - - # Instance seg head - instance_embed = F.conv2d(pixel_embed, seg_sd["instance_seg_head.weight"], - seg_sd["instance_seg_head.bias"]) - save_tensor(os.path.join(args.outdir, "seg_instance_embed"), instance_embed) - - # Mask predictor - obj_queries = normed_output.permute(1, 0, 2) # [1, NQ, D] - mask_embed = obj_queries - for j in range(3): - w = seg_sd[f"mask_predictor.mask_embed.layers.{j}.weight"] - b = seg_sd[f"mask_predictor.mask_embed.layers.{j}.bias"] - mask_embed = F.linear(mask_embed, w, b) - if j < 2: - mask_embed = F.relu(mask_embed) - save_tensor(os.path.join(args.outdir, "seg_mask_embed"), mask_embed) - - # Einsum: bqc,bchw -> bqhw - mask_preds = torch.einsum("bqc,bchw->bqhw", mask_embed, instance_embed) - save_tensor(os.path.join(args.outdir, "seg_mask_logits"), mask_preds) - print(f" mask_preds shape: {list(mask_preds.shape)}") - - # ═══════════════════════════════════════════════════════════════════════ - # End-to-end summary - # ═══════════════════════════════════════════════════════════════════════ - print("\n=== End-to-end Summary ===") - with torch.no_grad(): - class_logits = scores[-1, :, :, 0] # [1, NQ] - pres_sig = pres_logit.squeeze(-1).squeeze(0).sigmoid() # scalar - probs = class_logits.sigmoid() * pres_sig - top_k = min(10, NQ) - top_probs, top_idx = probs[0].topk(top_k) - print(f" Top {top_k} detections:") - for i in range(top_k): - qi = top_idx[i].item() - box = reference_boxes[qi, 0] - print(f" q{qi}: score={top_probs[i]:.4f}, box=({box[0]:.3f},{box[1]:.3f},{box[2]:.3f},{box[3]:.3f})") - - print(f"\nDone! Reference tensors saved to {args.outdir}") - - -if __name__ == "__main__": - main() diff --git a/tests/dump_phase6_reference.py b/tests/dump_phase6_reference.py deleted file mode 100644 index e124125..0000000 --- a/tests/dump_phase6_reference.py +++ /dev/null @@ -1,493 +0,0 @@ -#!/usr/bin/env python3 -""" -Dump Phase 6 (SAM prompt encoder + SAM mask decoder) reference tensors for -numerical comparison against sam3.cpp. - -The script reuses tracker neck features from Phase 3 so Phase 6 can be audited -in isolation and the fixtures can be reused by later tracker stages. - -Usage: - uv run python tests/dump_phase6_reference.py \ - --checkpoint raw_weights/sam3.pt \ - --prephase-ref tests/ref_phase3 \ - --cases tests/phase6_cases.tsv \ - --outdir tests/ref_phase6 -""" - -import argparse -import math -import os -from dataclasses import dataclass -from typing import Dict, List - -import numpy as np -import torch -import torch.nn.functional as F - - -@dataclass -class PromptCase: - case_id: str - multimask: bool - pos_points: List[List[float]] - neg_points: List[List[float]] - box: List[float] - - -def load_tensor(path: str) -> torch.Tensor: - with open(path + ".shape", "r", encoding="utf-8") as f: - shape = [int(x) for x in f.read().strip().split(",") if x] - data = np.fromfile(path + ".bin", dtype=np.float32).reshape(shape) - return torch.from_numpy(data) - - -def save_raw(path: str, arr: np.ndarray, shape) -> None: - os.makedirs(os.path.dirname(path) or ".", exist_ok=True) - with open(path + ".bin", "wb") as f: - f.write(arr.astype(np.float32, copy=False).tobytes()) - with open(path + ".shape", "w", encoding="utf-8") as f: - f.write(",".join(str(int(d)) for d in shape)) - - -def save_ggml_bd(path: str, x_bd: torch.Tensor) -> None: - assert x_bd.ndim == 2 - x_bd = x_bd.detach().cpu().float().contiguous() - b, d = x_bd.shape - save_raw(path, x_bd.numpy(), (d, b)) - - -def save_ggml_bnd(path: str, x_bnd: torch.Tensor) -> None: - assert x_bnd.ndim == 3 - x_bnd = x_bnd.detach().cpu().float().contiguous() - b, n, d = x_bnd.shape - save_raw(path, x_bnd.numpy(), (d, n, b)) - - -def save_ggml_bhwc(path: str, x_bhwc: torch.Tensor) -> None: - assert x_bhwc.ndim == 4 - x_bhwc = x_bhwc.detach().cpu().float().contiguous() - b, h, w, c = x_bhwc.shape - save_raw(path, x_bhwc.numpy(), (c, w, h, b)) - - -def save_ggml_nchw(path: str, x_nchw: torch.Tensor) -> None: - assert x_nchw.ndim == 4 - save_ggml_bhwc(path, x_nchw.permute(0, 2, 3, 1)) - - -def save_ggml_whcb_from_nchw(path: str, x_nchw: torch.Tensor) -> None: - assert x_nchw.ndim == 4 - x_nchw = x_nchw.detach().cpu().float().contiguous() - b, c, h, w = x_nchw.shape - save_raw(path, x_nchw.numpy(), (w, h, c, b)) - - -def save_ggml_masks(path: str, masks_bchw: torch.Tensor) -> None: - assert masks_bchw.ndim == 4 - b, c, h, w = masks_bchw.shape - save_ggml_bnd(path, masks_bchw.reshape(b, c, h * w)) - - -def parse_points(field: str) -> List[List[float]]: - if not field: - return [] - points = [] - for part in field.split("|"): - x_str, y_str = part.split(":") - points.append([float(x_str), float(y_str)]) - return points - - -def parse_box(field: str) -> List[float]: - if not field: - return [] - x0, y0, x1, y1 = field.split(":") - return [float(x0), float(y0), float(x1), float(y1)] - - -def load_cases(path: str) -> List[PromptCase]: - cases: List[PromptCase] = [] - with open(path, "r", encoding="utf-8") as f: - for line in f: - line = line.rstrip("\n") - if not line: - continue - fields = line.split("\t") - while len(fields) < 5: - fields.append("") - cases.append( - PromptCase( - case_id=fields[0], - multimask=fields[1] == "1", - pos_points=parse_points(fields[2]), - neg_points=parse_points(fields[3]), - box=parse_box(fields[4]), - ) - ) - return cases - - -def build_point_inputs(case: PromptCase): - coords: List[List[float]] = [] - labels: List[int] = [] - - if case.box: - coords.append([case.box[0], case.box[1]]) - coords.append([case.box[2], case.box[3]]) - labels.extend([2, 3]) - - coords.extend(case.pos_points) - labels.extend([1] * len(case.pos_points)) - - coords.extend(case.neg_points) - labels.extend([0] * len(case.neg_points)) - - coords_t = torch.tensor([coords], dtype=torch.float32) if coords else torch.zeros(1, 0, 2) - labels_t = torch.tensor([labels], dtype=torch.int64) if labels else torch.zeros(1, 0, dtype=torch.int64) - return coords_t, labels_t - - -def position_encode(coords_norm: torch.Tensor, pe_gaussian: torch.Tensor) -> torch.Tensor: - coords = 2.0 * coords_norm - 1.0 - enc = coords @ pe_gaussian - enc = 2.0 * math.pi * enc - return torch.cat([torch.sin(enc), torch.cos(enc)], dim=-1) - - -def build_prompt_encoder_outputs(weights: Dict[str, torch.Tensor], case: PromptCase): - embed_dim = 256 - input_size = 1008 - feat_size = 72 - - pe_gaussian = weights["pe_layer.positional_encoding_gaussian_matrix"].float() - coords, labels = build_point_inputs(case) - - pad_point = torch.zeros((coords.shape[0], 1, 2), dtype=torch.float32) - pad_label = -torch.ones((labels.shape[0], 1), dtype=torch.int64) - coords = torch.cat([coords, pad_point], dim=1) - labels = torch.cat([labels, pad_label], dim=1) - - coords_centered = coords + 0.5 - coords_norm = coords_centered.clone() - coords_norm[:, :, 0] /= input_size - coords_norm[:, :, 1] /= input_size - point_embedding = position_encode(coords_norm, pe_gaussian) - - not_a_point = weights["not_a_point_embed.weight"].float()[0] - point_embeds = [ - weights[f"point_embeddings.{i}.weight"].float()[0] - for i in range(4) - ] - - sparse = point_embedding.clone() - for idx in range(labels.shape[1]): - label = int(labels[0, idx].item()) - if label == -1: - sparse[:, idx, :] = not_a_point - else: - sparse[:, idx, :] = sparse[:, idx, :] + point_embeds[label] - - ys = (torch.arange(feat_size, dtype=torch.float32) + 0.5) / feat_size - xs = (torch.arange(feat_size, dtype=torch.float32) + 0.5) / feat_size - yy, xx = torch.meshgrid(ys, xs, indexing="ij") - dense_coords = torch.stack([xx, yy], dim=-1) - image_pe = position_encode(dense_coords, pe_gaussian).permute(2, 0, 1).unsqueeze(0) - - no_mask = weights["no_mask_embed.weight"].float().view(1, embed_dim, 1, 1) - dense = no_mask.expand(1, embed_dim, feat_size, feat_size).contiguous() - return sparse, dense, image_pe - - -def attention_forward(q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - prefix: str, - weights: Dict[str, torch.Tensor], - num_heads: int) -> torch.Tensor: - q_w = weights[prefix + ".q_proj.weight"].float() - q_b = weights[prefix + ".q_proj.bias"].float() - k_w = weights[prefix + ".k_proj.weight"].float() - k_b = weights[prefix + ".k_proj.bias"].float() - v_w = weights[prefix + ".v_proj.weight"].float() - v_b = weights[prefix + ".v_proj.bias"].float() - out_w = weights[prefix + ".out_proj.weight"].float() - out_b = weights[prefix + ".out_proj.bias"].float() - - q_proj = F.linear(q, q_w, q_b) - k_proj = F.linear(k, k_w, k_b) - v_proj = F.linear(v, v_w, v_b) - - bsz, nq, dim = q_proj.shape - nk = k_proj.shape[1] - head_dim = dim // num_heads - - qh = q_proj.view(bsz, nq, num_heads, head_dim).permute(0, 2, 1, 3) - kh = k_proj.view(bsz, nk, num_heads, head_dim).permute(0, 2, 1, 3) - vh = v_proj.view(bsz, nk, num_heads, head_dim).permute(0, 2, 1, 3) - - out = F.scaled_dot_product_attention(qh, kh, vh) - out = out.permute(0, 2, 1, 3).reshape(bsz, nq, dim) - return F.linear(out, out_w, out_b) - - -def mlp_forward(x: torch.Tensor, - weights: Dict[str, torch.Tensor], - prefix: str, - num_layers: int, - sigmoid_output: bool = False) -> torch.Tensor: - for i in range(num_layers): - w = weights[f"{prefix}.layers.{i}.weight"].float() - b = weights[f"{prefix}.layers.{i}.bias"].float() - x = F.linear(x, w, b) - if i < num_layers - 1: - x = F.relu(x) - if sigmoid_output: - x = torch.sigmoid(x) - return x - - -def layer_norm_2d(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: - u = x.mean(1, keepdim=True) - s = (x - u).pow(2).mean(1, keepdim=True) - x = (x - u) / torch.sqrt(s + eps) - return weight[:, None, None] * x + bias[:, None, None] - - -def dump_case(case: PromptCase, - prompt_weights: Dict[str, torch.Tensor], - mask_weights: Dict[str, torch.Tensor], - no_mem_embed: torch.Tensor, - neck_trk_0: torch.Tensor, - neck_trk_1: torch.Tensor, - neck_trk_2: torch.Tensor, - outdir: str) -> None: - os.makedirs(outdir, exist_ok=True) - - sparse_embeddings, dense_embeddings, image_pe = build_prompt_encoder_outputs(prompt_weights, case) - - image_embeddings = neck_trk_2 + no_mem_embed.view(1, -1, 1, 1) - feat_s0 = F.conv2d( - neck_trk_0, - mask_weights["conv_s0.weight"].float(), - mask_weights["conv_s0.bias"].float(), - ) - feat_s1 = F.conv2d( - neck_trk_1, - mask_weights["conv_s1.weight"].float(), - mask_weights["conv_s1.bias"].float(), - ) - - save_ggml_bnd(os.path.join(outdir, "sam_pe_sparse"), sparse_embeddings) - save_ggml_nchw(os.path.join(outdir, "sam_pe_dense"), dense_embeddings) - save_ggml_nchw(os.path.join(outdir, "sam_pe_image_pe"), image_pe) - save_ggml_nchw(os.path.join(outdir, "sam_dec_image_feats"), image_embeddings) - - output_tokens = torch.cat( - [ - mask_weights["obj_score_token.weight"].float(), - mask_weights["iou_token.weight"].float(), - mask_weights["mask_tokens.weight"].float(), - ], - dim=0, - ).unsqueeze(0).expand(sparse_embeddings.size(0), -1, -1) - tokens = torch.cat((output_tokens, sparse_embeddings), dim=1) - save_ggml_bnd(os.path.join(outdir, "sam_dec_tokens_initial"), tokens) - - src = image_embeddings + dense_embeddings - pos_src = image_pe.repeat(tokens.shape[0], 1, 1, 1) - bsz, dim, h, w = src.shape - keys = src.flatten(2).permute(0, 2, 1) - key_pe = pos_src.flatten(2).permute(0, 2, 1) - queries = tokens - query_pe = tokens - - for idx in range(2): - prefix = f"transformer.layers.{idx}" - if idx == 0: - queries = attention_forward( - queries, queries, queries, prefix + ".self_attn", mask_weights, 8 - ) - else: - q = queries + query_pe - attn_out = attention_forward(q, q, queries, prefix + ".self_attn", mask_weights, 8) - queries = queries + attn_out - queries = F.layer_norm( - queries, - [dim], - mask_weights[prefix + ".norm1.weight"].float(), - mask_weights[prefix + ".norm1.bias"].float(), - ) - - q = queries + query_pe - k = keys + key_pe - attn_out = attention_forward( - q, k, keys, prefix + ".cross_attn_token_to_image", mask_weights, 8 - ) - queries = F.layer_norm( - queries + attn_out, - [dim], - mask_weights[prefix + ".norm2.weight"].float(), - mask_weights[prefix + ".norm2.bias"].float(), - ) - - mlp_hidden = F.linear( - queries, - mask_weights[prefix + ".mlp.lin1.weight"].float(), - mask_weights[prefix + ".mlp.lin1.bias"].float(), - ) - mlp_hidden = F.relu(mlp_hidden) - mlp_hidden = F.linear( - mlp_hidden, - mask_weights[prefix + ".mlp.lin2.weight"].float(), - mask_weights[prefix + ".mlp.lin2.bias"].float(), - ) - queries = F.layer_norm( - queries + mlp_hidden, - [dim], - mask_weights[prefix + ".norm3.weight"].float(), - mask_weights[prefix + ".norm3.bias"].float(), - ) - - q = queries + query_pe - k = keys + key_pe - attn_out = attention_forward( - k, q, queries, prefix + ".cross_attn_image_to_token", mask_weights, 8 - ) - keys = F.layer_norm( - keys + attn_out, - [dim], - mask_weights[prefix + ".norm4.weight"].float(), - mask_weights[prefix + ".norm4.bias"].float(), - ) - - save_ggml_bnd(os.path.join(outdir, f"sam_dec_block{idx}_queries"), queries) - save_ggml_bnd(os.path.join(outdir, f"sam_dec_block{idx}_keys"), keys) - - q = queries + query_pe - k = keys + key_pe - attn_out = attention_forward( - q, k, keys, "transformer.final_attn_token_to_image", mask_weights, 8 - ) - queries = F.layer_norm( - queries + attn_out, - [dim], - mask_weights["transformer.norm_final_attn.weight"].float(), - mask_weights["transformer.norm_final_attn.bias"].float(), - ) - save_ggml_bnd(os.path.join(outdir, "sam_dec_final_queries"), queries) - - iou_token_out = queries[:, 1, :] - mask_tokens_out = queries[:, 2:6, :] - obj_token_out = queries[:, 0, :] - save_ggml_bnd(os.path.join(outdir, "sam_dec_mask_tokens"), mask_tokens_out) - save_ggml_bd(os.path.join(outdir, "sam_dec_sam_token"), mask_tokens_out[:, 0, :]) - - src_img = keys.transpose(1, 2).view(bsz, dim, h, w) - up1 = F.conv_transpose2d( - src_img, - mask_weights["output_upscaling.0.weight"].float(), - mask_weights["output_upscaling.0.bias"].float(), - stride=2, - ) - up1 = up1 + feat_s1 - up1 = layer_norm_2d( - up1, - mask_weights["output_upscaling.1.weight"].float(), - mask_weights["output_upscaling.1.bias"].float(), - ) - up1 = F.gelu(up1) - - up2 = F.conv_transpose2d( - up1, - mask_weights["output_upscaling.3.weight"].float(), - mask_weights["output_upscaling.3.bias"].float(), - stride=2, - ) - up2 = F.gelu(up2 + feat_s0) - - save_ggml_whcb_from_nchw(os.path.join(outdir, "sam_dec_feat_s1_proj"), feat_s1) - save_ggml_whcb_from_nchw(os.path.join(outdir, "sam_dec_feat_s0_proj"), feat_s0) - save_ggml_nchw(os.path.join(outdir, "sam_dec_upscaled"), up2) - - hyper_in_list = [ - mlp_forward( - mask_tokens_out[:, i, :], - mask_weights, - f"output_hypernetworks_mlps.{i}", - 3, - ) - for i in range(4) - ] - hyper_in = torch.stack(hyper_in_list, dim=1) - bsz, c, h, w = up2.shape - masks = (hyper_in @ up2.view(bsz, c, h * w)).view(bsz, -1, h, w) - iou_pred = mlp_forward( - iou_token_out, - mask_weights, - "iou_prediction_head", - 3, - sigmoid_output=True, - ) - obj_score_logits = mlp_forward( - obj_token_out, - mask_weights, - "pred_obj_score_head", - 3, - sigmoid_output=False, - ) - - save_ggml_masks(os.path.join(outdir, "sam_dec_masks"), masks) - save_ggml_bd(os.path.join(outdir, "sam_dec_iou"), iou_pred) - save_ggml_bd(os.path.join(outdir, "sam_dec_obj_score"), obj_score_logits) - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--checkpoint", required=True) - parser.add_argument("--prephase-ref", default="tests/ref_phase3") - parser.add_argument("--cases", default="tests/phase6_cases.tsv") - parser.add_argument("--outdir", default="tests/ref_phase6") - args = parser.parse_args() - - os.makedirs(args.outdir, exist_ok=True) - - print("Loading checkpoint...") - ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=False) - if "model" in ckpt: - ckpt = ckpt["model"] - - prompt_weights = { - k[len("tracker.sam_prompt_encoder."):]: v - for k, v in ckpt.items() - if k.startswith("tracker.sam_prompt_encoder.") - } - mask_weights = { - k[len("tracker.sam_mask_decoder."):]: v - for k, v in ckpt.items() - if k.startswith("tracker.sam_mask_decoder.") - } - no_mem_embed = ckpt["tracker.no_mem_embed"].float() - - neck_trk_0 = load_tensor(os.path.join(args.prephase_ref, "neck_trk_0")).float() - neck_trk_1 = load_tensor(os.path.join(args.prephase_ref, "neck_trk_1")).float() - neck_trk_2 = load_tensor(os.path.join(args.prephase_ref, "neck_trk_2")).float() - - cases = load_cases(args.cases) - for case in cases: - print(f"Dumping {case.case_id}...") - dump_case( - case, - prompt_weights, - mask_weights, - no_mem_embed, - neck_trk_0, - neck_trk_1, - neck_trk_2, - os.path.join(args.outdir, case.case_id), - ) - - print(f"All tensors saved to: {args.outdir}") - - -if __name__ == "__main__": - main() diff --git a/tests/dump_phase7_reference.py b/tests/dump_phase7_reference.py deleted file mode 100644 index 08a3c3c..0000000 --- a/tests/dump_phase7_reference.py +++ /dev/null @@ -1,836 +0,0 @@ -#!/usr/bin/env python3 -""" -Dump Phase 7 (video tracking) reference tensors for numerical comparison. - -This reuses Phase 3 tracker features and Phase 6 SAM outputs so later tracker work -can audit Phase 7 in isolation from earlier stages. - -Usage: - uv run python tests/dump_phase7_reference.py \ - --checkpoint raw_weights/sam3.pt \ - --prephase-ref tests/ref_phase3 \ - --phase6-ref tests/ref_phase6 \ - --cases tests/phase7_cases.tsv \ - --outdir tests/ref_phase7 -""" - -import argparse -import math -import os -from dataclasses import dataclass -from typing import Dict, List - -import numpy as np -import torch -import torch.nn.functional as F - - -@dataclass -class Phase7Case: - case_id: str - label: str - mem_cases: List[str] - spatial_tpos: List[int] - ptr_tpos: List[int] - - -def load_shape(path: str) -> List[int]: - with open(path + ".shape", "r", encoding="utf-8") as f: - return [int(x) for x in f.read().strip().split(",") if x] - - -def load_tensor(path: str) -> torch.Tensor: - shape = load_shape(path) - data = np.fromfile(path + ".bin", dtype=np.float32).reshape(shape) - return torch.from_numpy(data) - - -def load_ggml_bd(path: str) -> torch.Tensor: - d, b = load_shape(path) - data = np.fromfile(path + ".bin", dtype=np.float32).reshape((b, d)) - return torch.from_numpy(data) - - -def load_ggml_bnd(path: str) -> torch.Tensor: - d, n, b = load_shape(path) - data = np.fromfile(path + ".bin", dtype=np.float32).reshape((b, n, d)) - return torch.from_numpy(data) - - -def load_ggml_masks(path: str, h: int, w: int) -> torch.Tensor: - x = load_ggml_bnd(path) - b, n, hw = x.shape - assert hw == h * w - return x.view(b, n, h, w) - - -def save_raw(path: str, arr: np.ndarray, shape) -> None: - os.makedirs(os.path.dirname(path) or ".", exist_ok=True) - with open(path + ".bin", "wb") as f: - f.write(arr.astype(np.float32, copy=False).tobytes()) - with open(path + ".shape", "w", encoding="utf-8") as f: - f.write(",".join(str(int(d)) for d in shape)) - - -def save_tensor(path: str, t: torch.Tensor) -> None: - t = t.detach().cpu().float().contiguous() - save_raw(path, t.numpy(), t.shape) - - -def save_ggml_bd(path: str, x_bd: torch.Tensor) -> None: - assert x_bd.ndim == 2 - x_bd = x_bd.detach().cpu().float().contiguous() - b, d = x_bd.shape - save_raw(path, x_bd.numpy(), (d, b)) - - -def save_ggml_bnd(path: str, x_bnd: torch.Tensor) -> None: - assert x_bnd.ndim == 3 - x_bnd = x_bnd.detach().cpu().float().contiguous() - b, n, d = x_bnd.shape - save_raw(path, x_bnd.numpy(), (d, n, b)) - - -def save_ggml_bhwc(path: str, x_bhwc: torch.Tensor) -> None: - assert x_bhwc.ndim == 4 - x_bhwc = x_bhwc.detach().cpu().float().contiguous() - b, h, w, c = x_bhwc.shape - save_raw(path, x_bhwc.numpy(), (c, w, h, b)) - - -def save_ggml_nchw(path: str, x_nchw: torch.Tensor) -> None: - assert x_nchw.ndim == 4 - save_ggml_bhwc(path, x_nchw.permute(0, 2, 3, 1)) - - -def save_ggml_masks(path: str, masks_bchw: torch.Tensor) -> None: - assert masks_bchw.ndim == 4 - b, c, h, w = masks_bchw.shape - save_ggml_bnd(path, masks_bchw.reshape(b, c, h * w)) - - -def parse_int_list(field: str) -> List[int]: - if not field: - return [] - vals = [] - for part in field.split("|"): - vals.append(-1 if part == "-" else int(part)) - return vals - - -def load_cases(path: str) -> List[Phase7Case]: - cases: List[Phase7Case] = [] - with open(path, "r", encoding="utf-8") as f: - for line in f: - line = line.rstrip("\n") - if not line: - continue - fields = line.split("\t") - while len(fields) < 5: - fields.append("") - mem_cases = fields[2].split("|") if fields[2] else [] - spatial_tpos = parse_int_list(fields[3]) - ptr_tpos = parse_int_list(fields[4]) - assert len(mem_cases) == len(spatial_tpos) == len(ptr_tpos) - cases.append( - Phase7Case( - case_id=fields[0], - label=fields[1], - mem_cases=mem_cases, - spatial_tpos=spatial_tpos, - ptr_tpos=ptr_tpos, - ) - ) - return cases - - -def layer_norm_2d(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: - u = x.mean(1, keepdim=True) - s = (x - u).pow(2).mean(1, keepdim=True) - x = (x - u) / torch.sqrt(s + eps) - return weight[:, None, None] * x + bias[:, None, None] - - -def sinusoidal_pe(h: int, w: int, d_model: int) -> torch.Tensor: - half = d_model // 2 - scale = 2.0 * math.pi - temperature = 10000.0 - - ys = torch.arange(1, h + 1, dtype=torch.float32).view(h, 1).repeat(1, w) - xs = torch.arange(1, w + 1, dtype=torch.float32).view(1, w).repeat(h, 1) - ys = ys / h * scale - xs = xs / w * scale - - dim_t = torch.arange(half, dtype=torch.float32) - dim_t = temperature ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / half) - - pos_x = xs[:, :, None] / dim_t - pos_y = ys[:, :, None] / dim_t - pos_x = torch.stack((pos_x[:, :, 0::2].sin(), pos_x[:, :, 1::2].cos()), dim=3).flatten(2) - pos_y = torch.stack((pos_y[:, :, 0::2].sin(), pos_y[:, :, 1::2].cos()), dim=3).flatten(2) - return torch.cat((pos_y, pos_x), dim=2).permute(2, 0, 1).unsqueeze(0).contiguous() - - -def get_1d_sine_pe(pos_inds: torch.Tensor, dim: int, temperature: float = 10000.0) -> torch.Tensor: - pe_dim = dim // 2 - dim_t = torch.arange(pe_dim, dtype=torch.float32, device=pos_inds.device) - dim_t = temperature ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / pe_dim) - pos_embed = pos_inds.unsqueeze(-1) / dim_t - return torch.cat([pos_embed.sin(), pos_embed.cos()], dim=-1) - - -def init_t_xy(end_x, end_y, scale=1.0, offset=0, device=None): - t = torch.arange(end_x * end_y, dtype=torch.float32, device=device) - t_x = (t % end_x).float() - t_y = torch.div(t, end_x, rounding_mode="floor").float() - return t_x * scale + offset, t_y * scale + offset - - -def compute_axial_cis(dim, end_x, end_y, theta=10000.0, scale_pos=1.0, offset=0, device=None): - freqs_x = 1.0 / (theta ** (torch.arange(0, dim, 4, device=device)[: (dim // 4)].float() / dim)) - freqs_y = 1.0 / (theta ** (torch.arange(0, dim, 4, device=device)[: (dim // 4)].float() / dim)) - t_x, t_y = init_t_xy(end_x, end_y, scale_pos, offset, device=device) - freqs_x = torch.outer(t_x, freqs_x) - freqs_y = torch.outer(t_y, freqs_y) - freqs_cis_x = torch.polar(torch.ones_like(freqs_x), freqs_x) - freqs_cis_y = torch.polar(torch.ones_like(freqs_y), freqs_y) - return torch.cat([freqs_cis_x, freqs_cis_y], dim=-1) - - -def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor) -> torch.Tensor: - shape = [d if i >= x.ndim - 2 else 1 for i, d in enumerate(x.shape)] - return freqs_cis.view(*shape) - - -def apply_rotary_enc(xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor, repeat_freqs_k: bool = False): - xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) - xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2)) if xk.shape[-2] != 0 else None - freqs_cis = reshape_for_broadcast(freqs_cis, xq_) - xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3) - if xk_ is None: - return xq_out.type_as(xq), xk - if repeat_freqs_k: - r = xk_.shape[-2] // xq_.shape[-2] - freqs_cis = freqs_cis.repeat(*([1] * (freqs_cis.ndim - 2)), r, 1) - xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3) - return xq_out.type_as(xq), xk_out.type_as(xk) - - -def attention_forward( - q_in: torch.Tensor, - k_in: torch.Tensor, - v_in: torch.Tensor, - prefix: str, - weights: Dict[str, torch.Tensor], - num_heads: int, - rope_cis: torch.Tensor | None = None, - repeat_freqs_k: bool = False, - num_k_exclude_rope: int = 0, -) -> torch.Tensor: - q = F.linear(q_in, weights[prefix + ".q_proj.weight"].float(), weights[prefix + ".q_proj.bias"].float()) - k = F.linear(k_in, weights[prefix + ".k_proj.weight"].float(), weights[prefix + ".k_proj.bias"].float()) - v = F.linear(v_in, weights[prefix + ".v_proj.weight"].float(), weights[prefix + ".v_proj.bias"].float()) - - bsz, nq, dim = q.shape - nk = k.shape[1] - head_dim = dim // num_heads - - qh = q.view(bsz, nq, num_heads, head_dim).permute(0, 2, 1, 3) - kh = k.view(bsz, nk, num_heads, head_dim).permute(0, 2, 1, 3) - vh = v.view(bsz, nk, num_heads, head_dim).permute(0, 2, 1, 3) - - if rope_cis is not None: - nk_rope = kh.shape[2] - num_k_exclude_rope - if nk_rope > 0: - qh, kh_rope = apply_rotary_enc( - qh, - kh[:, :, :nk_rope, :], - rope_cis, - repeat_freqs_k=repeat_freqs_k, - ) - if nk_rope < kh.shape[2]: - kh = torch.cat([kh_rope, kh[:, :, nk_rope:, :]], dim=2) - else: - kh = kh_rope - - out = F.scaled_dot_product_attention(qh, kh, vh) - out = out.permute(0, 2, 1, 3).reshape(bsz, nq, dim) - return F.linear(out, weights[prefix + ".out_proj.weight"].float(), weights[prefix + ".out_proj.bias"].float()) - - -def mlp_forward(x: torch.Tensor, weights: Dict[str, torch.Tensor], prefix: str, num_layers: int) -> torch.Tensor: - for i in range(num_layers): - w = weights[f"{prefix}.layers.{i}.weight"].float() - b = weights[f"{prefix}.layers.{i}.bias"].float() - x = F.linear(x, w, b) - if i < num_layers - 1: - x = F.relu(x) - return x - - -def build_no_point_prompt_encoder_outputs(prompt_weights: Dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - embed_dim = 256 - feat_size = 72 - - pe_gaussian = prompt_weights["pe_layer.positional_encoding_gaussian_matrix"].float() - coords = torch.zeros((1, 1, 2), dtype=torch.float32) - labels = -torch.ones((1, 1), dtype=torch.int64) - - coords_centered = coords + 0.5 - coords_norm = coords_centered.clone() - coords_norm[:, :, 0] /= 1008.0 - coords_norm[:, :, 1] /= 1008.0 - coords_enc = 2.0 * coords_norm - 1.0 - coords_enc = coords_enc @ pe_gaussian - coords_enc = 2.0 * math.pi * coords_enc - point_embedding = torch.cat([torch.sin(coords_enc), torch.cos(coords_enc)], dim=-1) - - not_a_point = prompt_weights["not_a_point_embed.weight"].float()[0] - sparse = point_embedding.clone() - sparse[:, 0, :] = not_a_point - - ys = (torch.arange(feat_size, dtype=torch.float32) + 0.5) / feat_size - xs = (torch.arange(feat_size, dtype=torch.float32) + 0.5) / feat_size - yy, xx = torch.meshgrid(ys, xs, indexing="ij") - dense_coords = torch.stack([xx, yy], dim=-1) - dense_coords = 2.0 * dense_coords - 1.0 - dense_pe = dense_coords @ pe_gaussian - dense_pe = 2.0 * math.pi * dense_pe - image_pe = torch.cat([torch.sin(dense_pe), torch.cos(dense_pe)], dim=-1).permute(2, 0, 1).unsqueeze(0) - - no_mask = prompt_weights["no_mask_embed.weight"].float().view(1, embed_dim, 1, 1) - dense = no_mask.expand(1, embed_dim, feat_size, feat_size).contiguous() - return sparse, dense, image_pe - - -def run_sam_decoder( - mask_weights: Dict[str, torch.Tensor], - prompt_weights: Dict[str, torch.Tensor], - image_embeddings: torch.Tensor, - neck_trk_0: torch.Tensor, - neck_trk_1: torch.Tensor, -) -> dict[str, torch.Tensor]: - sparse_embeddings, dense_embeddings, image_pe = build_no_point_prompt_encoder_outputs(prompt_weights) - - feat_s0 = F.conv2d( - neck_trk_0, - mask_weights["conv_s0.weight"].float(), - mask_weights["conv_s0.bias"].float(), - ) - feat_s1 = F.conv2d( - neck_trk_1, - mask_weights["conv_s1.weight"].float(), - mask_weights["conv_s1.bias"].float(), - ) - - output_tokens = torch.cat( - [ - mask_weights["obj_score_token.weight"].float(), - mask_weights["iou_token.weight"].float(), - mask_weights["mask_tokens.weight"].float(), - ], - dim=0, - ).unsqueeze(0) - tokens = torch.cat((output_tokens, sparse_embeddings), dim=1) - - src = image_embeddings + dense_embeddings - pos_src = image_pe - bsz, dim, h, w = src.shape - keys = src.flatten(2).permute(0, 2, 1) - key_pe = pos_src.flatten(2).permute(0, 2, 1) - queries = tokens - query_pe = tokens - - for idx in range(2): - prefix = f"transformer.layers.{idx}" - if idx == 0: - queries = attention_forward( - queries, queries, queries, prefix + ".self_attn", mask_weights, 8 - ) - else: - q = queries + query_pe - attn_out = attention_forward(q, q, queries, prefix + ".self_attn", mask_weights, 8) - queries = queries + attn_out - queries = F.layer_norm( - queries, - [dim], - mask_weights[prefix + ".norm1.weight"].float(), - mask_weights[prefix + ".norm1.bias"].float(), - ) - - q = queries + query_pe - k = keys + key_pe - attn_out = attention_forward( - q, k, keys, prefix + ".cross_attn_token_to_image", mask_weights, 8 - ) - queries = F.layer_norm( - queries + attn_out, - [dim], - mask_weights[prefix + ".norm2.weight"].float(), - mask_weights[prefix + ".norm2.bias"].float(), - ) - - mlp_hidden = F.linear( - queries, - mask_weights[prefix + ".mlp.lin1.weight"].float(), - mask_weights[prefix + ".mlp.lin1.bias"].float(), - ) - mlp_hidden = F.relu(mlp_hidden) - mlp_hidden = F.linear( - mlp_hidden, - mask_weights[prefix + ".mlp.lin2.weight"].float(), - mask_weights[prefix + ".mlp.lin2.bias"].float(), - ) - queries = F.layer_norm( - queries + mlp_hidden, - [dim], - mask_weights[prefix + ".norm3.weight"].float(), - mask_weights[prefix + ".norm3.bias"].float(), - ) - - q = queries + query_pe - k = keys + key_pe - attn_out = attention_forward( - k, q, queries, prefix + ".cross_attn_image_to_token", mask_weights, 8 - ) - keys = F.layer_norm( - keys + attn_out, - [dim], - mask_weights[prefix + ".norm4.weight"].float(), - mask_weights[prefix + ".norm4.bias"].float(), - ) - - q = queries + query_pe - k = keys + key_pe - attn_out = attention_forward( - q, k, keys, "transformer.final_attn_token_to_image", mask_weights, 8 - ) - queries = F.layer_norm( - queries + attn_out, - [dim], - mask_weights["transformer.norm_final_attn.weight"].float(), - mask_weights["transformer.norm_final_attn.bias"].float(), - ) - - iou_token_out = queries[:, 1, :] - mask_tokens_out = queries[:, 2:6, :] - obj_token_out = queries[:, 0, :] - - src_img = keys.transpose(1, 2).view(bsz, dim, h, w) - up1 = F.conv_transpose2d( - src_img, - mask_weights["output_upscaling.0.weight"].float(), - mask_weights["output_upscaling.0.bias"].float(), - stride=2, - ) - up1 = up1 + feat_s1 - up1 = layer_norm_2d( - up1, - mask_weights["output_upscaling.1.weight"].float(), - mask_weights["output_upscaling.1.bias"].float(), - ) - up1 = F.gelu(up1) - - up2 = F.conv_transpose2d( - up1, - mask_weights["output_upscaling.3.weight"].float(), - mask_weights["output_upscaling.3.bias"].float(), - stride=2, - ) - up2 = F.gelu(up2 + feat_s0) - - hyper_in_list = [ - mlp_forward(mask_tokens_out[:, i, :], mask_weights, f"output_hypernetworks_mlps.{i}", 3) - for i in range(4) - ] - hyper_in = torch.stack(hyper_in_list, dim=1) - _, c, uh, uw = up2.shape - masks = (hyper_in @ up2.view(bsz, c, uh * uw)).view(bsz, -1, uh, uw) - iou_pred = mlp_forward(iou_token_out, mask_weights, "iou_prediction_head", 3).sigmoid() - obj_score_logits = mlp_forward(obj_token_out, mask_weights, "pred_obj_score_head", 3) - - return { - "sparse": sparse_embeddings, - "dense": dense_embeddings, - "image_pe": image_pe, - "mask_logits": masks[:, :1, :, :].contiguous(), - "iou": iou_pred[:, :1].contiguous(), - "obj_score": obj_score_logits.contiguous(), - "sam_token": mask_tokens_out[:, 0, :].contiguous(), - } - - -def run_memory_encoder( - tracker_weights: Dict[str, torch.Tensor], - neck_trk_2: torch.Tensor, - low_res_mask_logits: torch.Tensor, - obj_score_logit: torch.Tensor, -) -> dict[str, torch.Tensor]: - high_res_mask = F.interpolate( - low_res_mask_logits.float(), - size=(1008, 1008), - mode="bilinear", - align_corners=False, - ) - mask_for_mem = torch.sigmoid(high_res_mask) - mask_for_mem = mask_for_mem * 20.0 - 10.0 - - x = F.interpolate( - mask_for_mem.float(), - size=(1152, 1152), - mode="bilinear", - align_corners=False, - antialias=True, - ) - - conv_ids = [0, 3, 6, 9] - norm_ids = [1, 4, 7, 10] - for conv_id, norm_id in zip(conv_ids, norm_ids): - x = F.conv2d( - x, - tracker_weights[f"maskmem_backbone.mask_downsampler.encoder.{conv_id}.weight"].float(), - tracker_weights[f"maskmem_backbone.mask_downsampler.encoder.{conv_id}.bias"].float(), - stride=2, - padding=1, - ) - x = layer_norm_2d( - x, - tracker_weights[f"maskmem_backbone.mask_downsampler.encoder.{norm_id}.weight"].float(), - tracker_weights[f"maskmem_backbone.mask_downsampler.encoder.{norm_id}.bias"].float(), - ) - x = F.gelu(x) - - x = F.conv2d( - x, - tracker_weights["maskmem_backbone.mask_downsampler.encoder.12.weight"].float(), - tracker_weights["maskmem_backbone.mask_downsampler.encoder.12.bias"].float(), - ) - - pix_proj = F.conv2d( - neck_trk_2, - tracker_weights["maskmem_backbone.pix_feat_proj.weight"].float(), - tracker_weights["maskmem_backbone.pix_feat_proj.bias"].float(), - ) - fused = pix_proj + x - - block_outs = [] - h = fused - for idx in range(2): - prefix = f"maskmem_backbone.fuser.layers.{idx}" - residual = h - h = F.conv2d( - h, - tracker_weights[prefix + ".dwconv.weight"].float(), - tracker_weights[prefix + ".dwconv.bias"].float(), - padding=3, - groups=h.shape[1], - ) - h = layer_norm_2d( - h, - tracker_weights[prefix + ".norm.weight"].float(), - tracker_weights[prefix + ".norm.bias"].float(), - ) - h = h.permute(0, 2, 3, 1) - h = F.linear( - h, - tracker_weights[prefix + ".pwconv1.weight"].float(), - tracker_weights[prefix + ".pwconv1.bias"].float(), - ) - h = F.gelu(h) - h = F.linear( - h, - tracker_weights[prefix + ".pwconv2.weight"].float(), - tracker_weights[prefix + ".pwconv2.bias"].float(), - ) - gamma = tracker_weights[prefix + ".gamma"].float().view(1, 1, 1, -1) - h = residual + (gamma * h).permute(0, 3, 1, 2) - block_outs.append(h) - - out = F.conv2d( - h, - tracker_weights["maskmem_backbone.out_proj.weight"].float(), - tracker_weights["maskmem_backbone.out_proj.bias"].float(), - ) - - is_obj_appearing = (obj_score_logit > 0).float().view(-1, 1, 1, 1) - no_obj_embed = tracker_weights["no_obj_embed_spatial"].float().view(1, -1, 1, 1) - out = out + (1.0 - is_obj_appearing) * no_obj_embed - pos = sinusoidal_pe(72, 72, 64) - - return { - "high_res_mask": high_res_mask, - "mask_scaled": mask_for_mem, - "mask_downsampled": x, - "pix_proj": pix_proj, - "fused_input": fused, - "fuser0": block_outs[0], - "fuser1": block_outs[1], - "output": out, - "pos": pos, - } - - -def run_obj_ptr( - tracker_weights: Dict[str, torch.Tensor], - sam_token: torch.Tensor, - obj_score_logit: torch.Tensor, -) -> torch.Tensor: - h = sam_token - for i in range(3): - h = F.linear( - h, - tracker_weights[f"obj_ptr_proj.layers.{i}.weight"].float(), - tracker_weights[f"obj_ptr_proj.layers.{i}.bias"].float(), - ) - if i < 2: - h = F.relu(h) - no_obj_ptr = tracker_weights["no_obj_ptr"].float() - is_obj = (obj_score_logit > 0).float() - return is_obj * h + (1.0 - is_obj) * no_obj_ptr - - -def build_prompt_and_pos( - tracker_weights: Dict[str, torch.Tensor], - mem_outputs: List[dict[str, torch.Tensor]], - obj_ptrs: List[torch.Tensor], - case: Phase7Case, -) -> tuple[torch.Tensor, torch.Tensor, int]: - prompt_tokens: List[torch.Tensor] = [] - prompt_pos: List[torch.Tensor] = [] - - for idx, mem_out in enumerate(mem_outputs): - feats = mem_out["output"].flatten(2).permute(2, 0, 1).contiguous() - pos = mem_out["pos"].flatten(2).permute(2, 0, 1).contiguous() - tpos = case.spatial_tpos[idx] - if tpos > 0: - add = tracker_weights["maskmem_tpos_enc"].float()[7 - tpos - 1].view(1, 1, -1) - pos = pos + add - prompt_tokens.append(feats) - prompt_pos.append(pos) - - num_obj_ptr_tokens = 0 - ptr_tokens_list: List[torch.Tensor] = [] - ptr_pos_list: List[torch.Tensor] = [] - for idx, obj_ptr in enumerate(obj_ptrs): - ptr_rel = case.ptr_tpos[idx] - if ptr_rel < 0: - continue - pos_ind = torch.tensor([ptr_rel / 15.0], dtype=torch.float32) - obj_pos = get_1d_sine_pe(pos_ind, dim=256) - obj_pos = F.linear( - obj_pos, - tracker_weights["obj_ptr_tpos_proj.weight"].float(), - tracker_weights["obj_ptr_tpos_proj.bias"].float(), - ) - ptr = obj_ptr.view(1, 1, 4, 64).permute(1, 2, 0, 3).reshape(4, 1, 64) - pos = obj_pos.view(1, 1, 64).repeat_interleave(4, dim=0) - ptr_tokens_list.append(ptr) - ptr_pos_list.append(pos) - - if ptr_tokens_list: - ptr_tokens = torch.cat(ptr_tokens_list, dim=0) - ptr_pos = torch.cat(ptr_pos_list, dim=0) - prompt_tokens.append(ptr_tokens) - prompt_pos.append(ptr_pos) - num_obj_ptr_tokens = ptr_tokens.shape[0] - - return torch.cat(prompt_tokens, dim=0), torch.cat(prompt_pos, dim=0), num_obj_ptr_tokens - - -def run_memory_attention( - tracker_weights: Dict[str, torch.Tensor], - curr_neck_trk_2: torch.Tensor, - prompt: torch.Tensor, - prompt_pos: torch.Tensor, - num_obj_ptr_tokens: int, -) -> dict[str, torch.Tensor]: - src = curr_neck_trk_2.flatten(2).permute(2, 0, 1).contiguous() - src_pos = sinusoidal_pe(72, 72, 256).flatten(2).permute(2, 0, 1).contiguous() - - x = src.permute(1, 0, 2).contiguous() + 0.1 * src_pos.permute(1, 0, 2).contiguous() - mem = prompt.permute(1, 0, 2).contiguous() - mem_pos = prompt_pos.permute(1, 0, 2).contiguous() - - rope_cis = compute_axial_cis(dim=256, end_x=72, end_y=72) - out = {"input": x} - for idx in range(4): - prefix = f"transformer.encoder.layers.{idx}" - - x_norm = F.layer_norm( - x, - [256], - tracker_weights[prefix + ".norm1.weight"].float(), - tracker_weights[prefix + ".norm1.bias"].float(), - ) - sa = attention_forward( - x_norm, - x_norm, - x_norm, - prefix + ".self_attn", - tracker_weights, - 1, - rope_cis=rope_cis, - ) - x = x + sa - out[f"layer{idx}_after_sa"] = x - - x_norm = F.layer_norm( - x, - [256], - tracker_weights[prefix + ".norm2.weight"].float(), - tracker_weights[prefix + ".norm2.bias"].float(), - ) - ca = attention_forward( - x_norm, - mem + mem_pos, - mem, - prefix + ".cross_attn_image", - tracker_weights, - 1, - rope_cis=rope_cis, - repeat_freqs_k=True, - num_k_exclude_rope=num_obj_ptr_tokens, - ) - x = x + ca - out[f"layer{idx}_after_ca"] = x - - x_norm = F.layer_norm( - x, - [256], - tracker_weights[prefix + ".norm3.weight"].float(), - tracker_weights[prefix + ".norm3.bias"].float(), - ) - ffn = F.linear( - x_norm, - tracker_weights[prefix + ".linear1.weight"].float(), - tracker_weights[prefix + ".linear1.bias"].float(), - ) - ffn = F.relu(ffn) - ffn = F.linear( - ffn, - tracker_weights[prefix + ".linear2.weight"].float(), - tracker_weights[prefix + ".linear2.bias"].float(), - ) - x = x + ffn - out[f"layer{idx}_after_ffn"] = x - - x = F.layer_norm( - x, - [256], - tracker_weights["transformer.encoder.norm.weight"].float(), - tracker_weights["transformer.encoder.norm.bias"].float(), - ) - out["output"] = x - return out - - -def write_meta(case_dir: str, case: Phase7Case) -> None: - with open(os.path.join(case_dir, "meta.txt"), "w", encoding="utf-8") as f: - f.write(f"num_slots={len(case.mem_cases)}\n") - for i, tpos in enumerate(case.spatial_tpos): - f.write(f"slot{i}_spatial_tpos={tpos}\n") - f.write(f"slot{i}_ptr_tpos={case.ptr_tpos[i]}\n") - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--checkpoint", required=True) - parser.add_argument("--prephase-ref", default="tests/ref_phase3") - parser.add_argument("--phase6-ref", default="tests/ref_phase6") - parser.add_argument("--cases", default="tests/phase7_cases.tsv") - parser.add_argument("--outdir", default="tests/ref_phase7") - args = parser.parse_args() - - os.makedirs(args.outdir, exist_ok=True) - - ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=False) - if "model" in ckpt: - ckpt = ckpt["model"] - - tracker_weights = { - k[len("tracker."):]: v - for k, v in ckpt.items() - if k.startswith("tracker.") - } - prompt_weights = { - k[len("tracker.sam_prompt_encoder."):]: v - for k, v in ckpt.items() - if k.startswith("tracker.sam_prompt_encoder.") - } - mask_weights = { - k[len("tracker.sam_mask_decoder."):]: v - for k, v in ckpt.items() - if k.startswith("tracker.sam_mask_decoder.") - } - - neck_trk_0 = load_tensor(os.path.join(args.prephase_ref, "neck_trk_0")).float() - neck_trk_1 = load_tensor(os.path.join(args.prephase_ref, "neck_trk_1")).float() - neck_trk_2 = load_tensor(os.path.join(args.prephase_ref, "neck_trk_2")).float() - - cases = load_cases(args.cases) - for case in cases: - case_dir = os.path.join(args.outdir, case.case_id) - os.makedirs(case_dir, exist_ok=True) - write_meta(case_dir, case) - - save_tensor(os.path.join(case_dir, "input_curr_neck_trk_0"), neck_trk_0) - save_tensor(os.path.join(case_dir, "input_curr_neck_trk_1"), neck_trk_1) - save_tensor(os.path.join(case_dir, "input_curr_neck_trk_2"), neck_trk_2) - - mem_outputs = [] - obj_ptrs = [] - for idx, mem_case in enumerate(case.mem_cases): - masks = load_ggml_masks(os.path.join(args.phase6_ref, mem_case, "sam_dec_masks"), 288, 288).float() - low_res_mask = masks[:, :1, :, :].contiguous() - sam_token = load_ggml_bd(os.path.join(args.phase6_ref, mem_case, "sam_dec_sam_token")).float() - obj_score = load_ggml_bd(os.path.join(args.phase6_ref, mem_case, "sam_dec_obj_score")).float() - - save_tensor(os.path.join(case_dir, f"input_mem_mask_logits_{idx}"), low_res_mask) - save_tensor(os.path.join(case_dir, f"input_mem_sam_token_{idx}"), sam_token) - save_tensor(os.path.join(case_dir, f"input_mem_obj_score_{idx}"), obj_score) - - mem_out = run_memory_encoder(tracker_weights, neck_trk_2, low_res_mask, obj_score) - obj_ptr = run_obj_ptr(tracker_weights, sam_token, obj_score) - mem_outputs.append(mem_out) - obj_ptrs.append(obj_ptr) - - save_ggml_nchw(os.path.join(case_dir, f"phase7_mem{idx}_high_res_mask"), mem_out["high_res_mask"]) - save_ggml_nchw(os.path.join(case_dir, f"phase7_mem{idx}_mask_scaled"), mem_out["mask_scaled"]) - save_ggml_nchw(os.path.join(case_dir, f"phase7_mem{idx}_mask_downsampled"), mem_out["mask_downsampled"]) - save_ggml_nchw(os.path.join(case_dir, f"phase7_mem{idx}_pix_proj"), mem_out["pix_proj"]) - save_ggml_nchw(os.path.join(case_dir, f"phase7_mem{idx}_fused_input"), mem_out["fused_input"]) - save_ggml_nchw(os.path.join(case_dir, f"phase7_mem{idx}_fuser0"), mem_out["fuser0"]) - save_ggml_nchw(os.path.join(case_dir, f"phase7_mem{idx}_fuser1"), mem_out["fuser1"]) - save_ggml_nchw(os.path.join(case_dir, f"phase7_mem{idx}_output"), mem_out["output"]) - save_ggml_bd(os.path.join(case_dir, f"phase7_obj_ptr{idx}"), obj_ptr) - - prompt, prompt_pos, num_obj_ptr_tokens = build_prompt_and_pos( - tracker_weights, mem_outputs, obj_ptrs, case - ) - mem_attn = run_memory_attention( - tracker_weights, - neck_trk_2, - prompt, - prompt_pos, - num_obj_ptr_tokens, - ) - - save_ggml_bnd(os.path.join(case_dir, "phase7_mem_attn_input"), mem_attn["input"]) - for idx in range(4): - save_ggml_bnd(os.path.join(case_dir, f"phase7_mem_attn_layer{idx}_after_sa"), mem_attn[f"layer{idx}_after_sa"]) - save_ggml_bnd(os.path.join(case_dir, f"phase7_mem_attn_layer{idx}_after_ca"), mem_attn[f"layer{idx}_after_ca"]) - save_ggml_bnd(os.path.join(case_dir, f"phase7_mem_attn_layer{idx}_after_ffn"), mem_attn[f"layer{idx}_after_ffn"]) - save_ggml_bnd(os.path.join(case_dir, "phase7_mem_attn_output"), mem_attn["output"]) - - conditioned = mem_attn["output"].permute(0, 2, 1).reshape(1, 256, 72, 72).contiguous() - prop = run_sam_decoder(mask_weights, prompt_weights, conditioned, neck_trk_0, neck_trk_1) - save_ggml_masks(os.path.join(case_dir, "phase7_prop_masks"), prop["mask_logits"]) - save_ggml_bd(os.path.join(case_dir, "phase7_prop_iou"), prop["iou"]) - save_ggml_bd(os.path.join(case_dir, "phase7_prop_obj_score"), prop["obj_score"]) - save_ggml_bd(os.path.join(case_dir, "phase7_prop_sam_token"), prop["sam_token"]) - - print(f"Saved Phase 7 references to {args.outdir}") - - -if __name__ == "__main__": - main() diff --git a/tests/dump_reference.py b/tests/dump_reference.py deleted file mode 100644 index 85b57f2..0000000 --- a/tests/dump_reference.py +++ /dev/null @@ -1,169 +0,0 @@ -#!/usr/bin/env python3 -""" -Dump reference intermediate tensors from the SAM3 Python model for C++ verification. -Saves .bin files with raw float32 data and .shape files with tensor dimensions. - -Usage: - python tests/dump_reference.py --checkpoint raw_weights/sam3.pt --image tests/test_truck.jpg --outdir tests/ref/ -""" -import argparse -import os -import sys -import struct -import math - -import numpy as np -import torch -import torch.nn.functional as F -from PIL import Image -from torchvision.transforms import v2 - -# Add SAM3 repo to path -sys.path.insert(0, os.path.expanduser("~/Documents/sam3")) -from sam3.model.vitdet import ViT, compute_axial_cis, window_partition, window_unpartition -from sam3.model_builder import _create_vit_backbone - - -def save_tensor(path, t): - """Save a tensor as raw float32 binary + shape file.""" - t = t.detach().cpu().float().contiguous() - with open(path + ".bin", "wb") as f: - f.write(t.numpy().tobytes()) - with open(path + ".shape", "w") as f: - f.write(",".join(str(d) for d in t.shape)) - print(f" saved {path} shape={list(t.shape)} bytes={t.numel()*4}") - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--checkpoint", required=True, help="Path to sam3.pt") - parser.add_argument("--image", required=True, help="Path to test image") - parser.add_argument("--outdir", default="tests/ref", help="Output directory") - args = parser.parse_args() - - os.makedirs(args.outdir, exist_ok=True) - device = "cpu" # Use CPU for reproducibility - - # ── Load model ───────────────────────────────────────────────────────── - print("Loading checkpoint...") - ckpt = torch.load(args.checkpoint, map_location=device, weights_only=False) - state_dict = ckpt if isinstance(ckpt, dict) and "model" not in ckpt else ckpt.get("model", ckpt) - - # Extract ViT weights - vit_prefix = "detector.backbone.visual.trunk." - vit_state = {k[len(vit_prefix):]: v for k, v in state_dict.items() if k.startswith(vit_prefix)} - - print(f" Found {len(vit_state)} ViT parameters") - - # ── Create ViT backbone ──────────────────────────────────────────────── - print("Creating ViT backbone...") - vit = _create_vit_backbone(compile_mode=None) - vit.use_act_checkpoint = False # Disable for inference - missing, unexpected = vit.load_state_dict(vit_state, strict=False) - if missing: - print(f" WARNING: Missing keys: {missing[:5]}...") - if unexpected: - print(f" WARNING: Unexpected keys: {unexpected[:5]}...") - vit.eval() - vit.to(device) - - # ── Load and preprocess image ────────────────────────────────────────── - print(f"Loading image: {args.image}") - img = Image.open(args.image).convert("RGB") - print(f" Original size: {img.size}") - - transform = v2.Compose([ - v2.ToDtype(torch.uint8, scale=True), - v2.Resize(size=(1008, 1008)), - v2.ToDtype(torch.float32, scale=True), - v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), - ]) - - img_tensor = v2.functional.to_image(img) - img_preprocessed = transform(img_tensor).unsqueeze(0).to(device) - print(f" Preprocessed shape: {list(img_preprocessed.shape)}") - save_tensor(os.path.join(args.outdir, "preprocessed"), img_preprocessed) - - # ── Step-by-step forward pass ────────────────────────────────────────── - with torch.no_grad(): - # 1. Patch embedding - x = vit.patch_embed(img_preprocessed) - print(f" After patch_embed: {list(x.shape)}") # [1, 72, 72, 1024] - save_tensor(os.path.join(args.outdir, "patch_embed"), x) - - # 2. Positional embedding - h, w = x.shape[1], x.shape[2] - from sam3.model.vitdet import get_abs_pos - pos = get_abs_pos(vit.pos_embed, vit.pretrain_use_cls_token, (h, w), - retain_cls_token=False, tiling=vit.tile_abs_pos) - print(f" Pos embed (tiled): {list(pos.shape)}") # [1, 72, 72, 1024] - save_tensor(os.path.join(args.outdir, "pos_embed_tiled"), pos) - - x = x + pos - save_tensor(os.path.join(args.outdir, "after_pos_embed"), x) - - # 3. ln_pre - x = vit.ln_pre(x) - save_tensor(os.path.join(args.outdir, "after_ln_pre"), x) - - # 4. Block 0 output - x = vit.blocks[0](x) - save_tensor(os.path.join(args.outdir, "block_0_out"), x) - - # 5. All remaining blocks - for i in range(1, len(vit.blocks)): - x = vit.blocks[i](x) - - save_tensor(os.path.join(args.outdir, "vit_final"), x) - - # 6. Permute to [B, C, H, W] - feats = x.permute(0, 3, 1, 2) - save_tensor(os.path.join(args.outdir, "vit_output_bchw"), feats) - - # ── RoPE reference ───────────────────────────────────────────────────── - print("\nDumping RoPE reference values...") - - # Window attention RoPE (24x24, rope_pt_size=24, no interp needed since pt_size == input_size) - # But wait — for window blocks, input_size is (24, 24), rope_pt_size is from the config - # Looking at the model_builder: rope_pt_size=None → defaults to (window_size, window_size) = (24, 24) - # And rope_interp=True, scale_pos = 24/24 = 1.0 → no scaling - freqs_cis_window = compute_axial_cis(dim=64, end_x=24, end_y=24, theta=10000.0, scale_pos=1.0) - save_tensor(os.path.join(args.outdir, "rope_window_complex"), freqs_cis_window) - # Convert to real (cos, sin) pairs: freqs_cis is complex [576, 32] - freqs_real_window = torch.view_as_real(freqs_cis_window) # [576, 32, 2] - save_tensor(os.path.join(args.outdir, "rope_window_real"), freqs_real_window) - - # Global attention RoPE (72x72, rope_pt_size=24, rope_interp=True → scale_pos=24/72=1/3) - # For global blocks: input_size = (72, 72), rope_pt_size = (24, 24) (from None default) - # Wait — actually, looking at the code: - # Block.__init__: input_size=input_size if window_size == 0 else (window_size, window_size) - # For global blocks, window_size=0, so input_size=(72, 72) - # Attention.__init__: rope_pt_size from Block, which is (24,24) for rope_pt_size=None - # scale_pos = rope_pt_size[0] / input_size[0] = 24 / 72 = 1/3 - scale_pos_global = 24.0 / 72.0 - freqs_cis_global = compute_axial_cis(dim=64, end_x=72, end_y=72, theta=10000.0, scale_pos=scale_pos_global) - freqs_real_global = torch.view_as_real(freqs_cis_global) # [5184, 32, 2] - save_tensor(os.path.join(args.outdir, "rope_global_real"), freqs_real_global) - - # ── Neck reference ───────────────────────────────────────────────────── - print("\nDumping neck reference values...") - # Load neck weights - det_neck_prefix = "detector.backbone.visual.convs." - # We need to instantiate the neck too, but for now just save the ViT output - # The neck test will be done separately once we verify ViT outputs match - - # ── Sinusoidal PE reference ──────────────────────────────────────────── - print("\nDumping sinusoidal PE reference values...") - from sam3.model.position_encoding import PositionEmbeddingSine - pe = PositionEmbeddingSine(num_pos_feats=256, normalize=True, scale=None, temperature=10000) - - for H, W, name in [(288, 288, "pe_288"), (144, 144, "pe_144"), (72, 72, "pe_72"), (36, 36, "pe_36")]: - dummy = torch.zeros(1, 1, H, W) - pe_out = pe(dummy) # [1, 256, H, W] - save_tensor(os.path.join(args.outdir, name), pe_out) - - print("\nDone! Reference tensors saved to", args.outdir) - - -if __name__ == "__main__": - main() diff --git a/tests/dump_sam2_pvs_reference.py b/tests/dump_sam2_pvs_reference.py deleted file mode 100644 index 0f20fbe..0000000 --- a/tests/dump_sam2_pvs_reference.py +++ /dev/null @@ -1,132 +0,0 @@ -#!/usr/bin/env python3 -"""Dump SAM2 PVS decoder reference tensors for C++ comparison. - -Usage: - cd ~/Documents/sam2 - SAM2_DUMP_DIR=/tmp/debug_sam2_pvs \ - python /path/to/dump_sam2_pvs_reference.py \ - --checkpoint /path/to/sam2.1_hiera_base_plus.pt \ - --config configs/sam2.1/sam2.1_hiera_b+.yaml \ - --image /path/to/test_image.jpg \ - --point-x 600 --point-y 599 -""" -import argparse, os, sys, numpy as np -sys.path.insert(0, os.path.expanduser("~/Documents/sam2")) -import torch -torch.set_default_dtype(torch.float32) -from PIL import Image -from sam2.build_sam import build_sam2 -from sam2.sam2_image_predictor import SAM2ImagePredictor - -DUMP_DIR = os.environ.get("SAM2_DUMP_DIR", "/tmp/debug_sam2_pvs") - -def save(name, t): - os.makedirs(DUMP_DIR, exist_ok=True) - t = t.detach().cpu().float().contiguous() - path = os.path.join(DUMP_DIR, name) - with open(path + ".bin", "wb") as f: - f.write(t.numpy().tobytes()) - with open(path + ".shape", "w") as f: - f.write(",".join(str(d) for d in t.shape)) - print(f" {name}: {tuple(t.shape)}") - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--checkpoint", required=True) - parser.add_argument("--config", required=True) - parser.add_argument("--image", required=True) - parser.add_argument("--point-x", type=float, default=600) - parser.add_argument("--point-y", type=float, default=599) - args = parser.parse_args() - - model = build_sam2(args.config, args.checkpoint, device="cpu") - model = model.float().eval() - predictor = SAM2ImagePredictor(model) - - image = np.array(Image.open(args.image).convert("RGB")) - print(f"Image: {args.image} ({image.shape[1]}x{image.shape[0]})") - - # Set image and get features - with torch.no_grad(): - predictor.set_image(image) - - # Dump features stored after set_image - print("\n=== Features after set_image ===") - save("features_image_embed", predictor._features["image_embed"]) - for i, f in enumerate(predictor._features["high_res_feats"]): - save(f"features_high_res_{i}", f) - - # Run PVS prediction - point_coords = np.array([[args.point_x, args.point_y]], dtype=np.float32) - point_labels = np.array([1], dtype=np.int32) - - print(f"\n=== PVS with point ({args.point_x}, {args.point_y}) ===") - - # Manually trace the prediction path for dumps - with torch.no_grad(): - # Step 1: Coordinate normalization (matching _prep_prompts) - unnorm_coords = predictor._transforms.transform_coords( - torch.from_numpy(point_coords).unsqueeze(0).float(), - normalize=True, - orig_hw=(image.shape[0], image.shape[1]) - ) - labels = torch.from_numpy(point_labels).unsqueeze(0).int() - concat_points = (unnorm_coords, labels) - save("unnorm_coords", unnorm_coords.squeeze(0)) - print(f" unnorm_coords: {unnorm_coords.squeeze(0).numpy()}") - - # Step 2: Prompt encoder - sparse_emb, dense_emb = model.sam_prompt_encoder( - points=concat_points, boxes=None, masks=None - ) - save("sparse_embeddings", sparse_emb) - save("dense_embeddings", dense_emb) - - # Step 3: Get dense PE - image_pe = model.sam_prompt_encoder.get_dense_pe() - save("image_pe", image_pe) - - # Step 4: Prepare high-res features - high_res_features = [ - feat_level[0].unsqueeze(0) - for feat_level in predictor._features["high_res_feats"] - ] - for i, hrf in enumerate(high_res_features): - save(f"high_res_for_decoder_{i}", hrf) - - # Step 5: Run mask decoder - image_embed = predictor._features["image_embed"][0].unsqueeze(0) - save("image_embed_for_decoder", image_embed) - - low_res_masks, iou_predictions, _, _ = model.sam_mask_decoder( - image_embeddings=image_embed, - image_pe=image_pe, - sparse_prompt_embeddings=sparse_emb, - dense_prompt_embeddings=dense_emb, - multimask_output=True, - repeat_image=False, - high_res_features=high_res_features, - ) - save("low_res_masks", low_res_masks) - save("iou_predictions", iou_predictions) - print(f" iou_predictions: {iou_predictions.numpy()}") - print(f" masks range: [{low_res_masks.min():.4f}, {low_res_masks.max():.4f}]") - - # Also run through the full predict API for comparison - with torch.no_grad(): - masks, scores, logits = predictor.predict( - point_coords=point_coords, - point_labels=point_labels, - multimask_output=True, - normalize_coords=True, - ) - save("api_masks", torch.from_numpy(masks).float()) - save("api_scores", torch.from_numpy(scores).float()) - save("api_logits", torch.from_numpy(logits).float()) - print(f"\n API scores: {scores}") - print(f" API masks shape: {masks.shape}") - - print(f"\nDone. Tensors in {DUMP_DIR}/") - -if __name__ == "__main__": - main() diff --git a/tests/dump_sam2_reference.py b/tests/dump_sam2_reference.py deleted file mode 100644 index db216ff..0000000 --- a/tests/dump_sam2_reference.py +++ /dev/null @@ -1,231 +0,0 @@ -#!/usr/bin/env python3 -"""Dump SAM2 reference tensors for C++ comparison. - -Usage: - cd ~/Documents/sam2 - SAM2_DUMP_DIR=/tmp/debug_sam2_pipeline \ - python /path/to/dump_sam2_reference.py \ - --checkpoint /path/to/sam2.1_hiera_base_plus.pt \ - --config configs/sam2.1/sam2.1_hiera_b+.yaml \ - --image /path/to/test_image.jpg \ - --stage all -""" - -import argparse -import os -import sys -import numpy as np - -# Add SAM2 to path -sys.path.insert(0, os.path.expanduser("~/Documents/sam2")) - -import torch -torch.set_default_dtype(torch.float32) - -from PIL import Image -from sam2.build_sam import build_sam2 -from sam2.sam2_image_predictor import SAM2ImagePredictor - -# Import dump utilities -sys.path.insert(0, os.path.expanduser("~/Documents/sam2")) -from debug_pipeline_utils import save_tensor, DUMP_DIR - - -def dump_preprocessing(predictor, image_np): - """Stage 0: Dump preprocessed image tensor.""" - print("\n=== Stage 0: Preprocessing ===") - # Run the transforms manually - from sam2.utils.transforms import SAM2Transforms - transforms = predictor._transforms - - # Get the preprocessed image - from torchvision.transforms import ToTensor - img_tensor = ToTensor()(Image.fromarray(image_np)) # [3, H, W] float [0,1] - - # Resize - from torchvision.transforms.functional import resize - img_resized = resize(img_tensor, [predictor.model.image_size, predictor.model.image_size], - antialias=True) - - # Normalize (ImageNet) - mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1) - std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1) - img_normalized = (img_resized - mean) / std - - save_tensor("preprocessed_image", img_normalized.unsqueeze(0)) - print(f" preprocessed_image: {img_normalized.shape}") - - # Also save as CHW float for C++ consumption - chw = img_normalized.numpy().astype(np.float32) - with open(os.path.join(DUMP_DIR, "preprocessed.bin"), "wb") as f: - f.write(chw.tobytes()) - print(f" preprocessed.bin: CHW float32, shape {chw.shape}") - - return img_normalized.unsqueeze(0) - - -def dump_backbone(predictor, img_batch): - """Stage 1: Dump Hiera backbone outputs.""" - print("\n=== Stage 1: Hiera Backbone ===") - model = predictor.model - - # Run backbone trunk - with torch.no_grad(): - backbone_out = model.image_encoder.trunk(img_batch) - - # backbone_out is a list of intermediate features - # For Hiera, the output is typically from the trunk's forward which returns intermediates - print(f" Backbone output type: {type(backbone_out)}") - - if isinstance(backbone_out, (list, tuple)): - for i, feat in enumerate(backbone_out): - save_tensor(f"hiera_stage_{i}", feat) - print(f" hiera_stage_{i}: {feat.shape}") - elif isinstance(backbone_out, dict): - for k, v in backbone_out.items(): - if isinstance(v, torch.Tensor): - save_tensor(f"backbone_{k}", v) - print(f" backbone_{k}: {v.shape}") - - return backbone_out - - -def dump_fpn_neck(predictor, img_batch): - """Stage 2: Dump FPN neck outputs (backbone + neck together).""" - print("\n=== Stage 2: FPN Neck ===") - model = predictor.model - - with torch.no_grad(): - backbone_out = model.image_encoder(img_batch) - - # backbone_out is a dict with: - # "vision_features": src - # "vision_pos_enc": pos - # "backbone_fpn": features - for k, v in backbone_out.items(): - if isinstance(v, torch.Tensor): - save_tensor(f"imgenc_{k}", v) - print(f" imgenc_{k}: {v.shape}") - elif isinstance(v, (list, tuple)): - for i, item in enumerate(v): - if isinstance(item, torch.Tensor): - save_tensor(f"imgenc_{k}_{i}", item) - print(f" imgenc_{k}_{i}: {item.shape}") - - return backbone_out - - -def dump_image_features(predictor, image_np): - """Stage 3: Dump prepared image features (what set_image produces).""" - print("\n=== Stage 3: Image Features ===") - - with torch.no_grad(): - predictor.set_image(image_np) - - # Access cached features - features = predictor._features - for k, v in features.items(): - if isinstance(v, torch.Tensor): - save_tensor(f"features_{k}", v) - print(f" features_{k}: {v.shape}") - elif isinstance(v, (list, tuple)): - for i, item in enumerate(v): - if isinstance(item, torch.Tensor): - save_tensor(f"features_{k}_{i}", item) - print(f" features_{k}_{i}: {item.shape}") - - return features - - -def dump_pvs_inference(predictor, image_np, point_coords, point_labels): - """Stage 4-5: Dump prompt encoding + mask decoding.""" - print("\n=== Stage 4-5: PVS Inference ===") - - with torch.no_grad(): - if not predictor._is_image_set: - predictor.set_image(image_np) - - masks, scores, logits = predictor.predict( - point_coords=point_coords, - point_labels=point_labels, - multimask_output=True, - normalize_coords=True, - ) - - save_tensor("pvs_masks", torch.from_numpy(masks).float()) - save_tensor("pvs_scores", torch.from_numpy(scores).float()) - save_tensor("pvs_logits", torch.from_numpy(logits).float()) - print(f" pvs_masks: {masks.shape}") - print(f" pvs_scores: {scores.shape}, values: {scores}") - print(f" pvs_logits: {logits.shape}") - - return masks, scores, logits - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--checkpoint", required=True) - parser.add_argument("--config", required=True) - parser.add_argument("--image", required=True) - parser.add_argument("--stage", default="all", - help="all, preprocess, backbone, fpn, features, pvs") - parser.add_argument("--point-x", type=float, default=None) - parser.add_argument("--point-y", type=float, default=None) - args = parser.parse_args() - - os.makedirs(DUMP_DIR, exist_ok=True) - print(f"Dump directory: {DUMP_DIR}") - - # Load model - print(f"\nLoading model: {args.config}") - model = build_sam2( - config_file=args.config, - ckpt_path=args.checkpoint, - device="cpu", - ) - model = model.float() - model.eval() - - predictor = SAM2ImagePredictor(model) - - # Load image - image = np.array(Image.open(args.image).convert("RGB")) - print(f"Image: {args.image} ({image.shape[1]}x{image.shape[0]})") - - # Default point at center of image - px = args.point_x if args.point_x is not None else image.shape[1] / 2 - py = args.point_y if args.point_y is not None else image.shape[0] / 2 - point_coords = np.array([[px, py]], dtype=np.float32) - point_labels = np.array([1], dtype=np.int32) - print(f"Test point: ({px}, {py})") - - stages = args.stage.split(",") if args.stage != "all" else [ - "preprocess", "backbone", "fpn", "features", "pvs" - ] - - img_batch = None - - if "preprocess" in stages: - img_batch = dump_preprocessing(predictor, image) - - if "backbone" in stages: - if img_batch is None: - img_batch = dump_preprocessing(predictor, image) - dump_backbone(predictor, img_batch) - - if "fpn" in stages: - if img_batch is None: - img_batch = dump_preprocessing(predictor, image) - dump_fpn_neck(predictor, img_batch) - - if "features" in stages: - dump_image_features(predictor, image) - - if "pvs" in stages: - dump_pvs_inference(predictor, image, point_coords, point_labels) - - print(f"\nDone. Tensors saved to {DUMP_DIR}/") - - -if __name__ == "__main__": - main() diff --git a/tests/dump_vit_reference.py b/tests/dump_vit_reference.py deleted file mode 100644 index f88cea4..0000000 --- a/tests/dump_vit_reference.py +++ /dev/null @@ -1,307 +0,0 @@ -#!/usr/bin/env python3 -""" -Dump ViT backbone intermediate tensors from the SAM3 checkpoint for C++ verification. -Uses inlined functions to avoid heavy SAM3 dependencies. - -Usage: - uv run python tests/dump_vit_reference.py --checkpoint raw_weights/sam3.pt --image tests/test_random.jpg -""" -import argparse -import math -import os -import sys -from functools import partial -from typing import Tuple - -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F -from PIL import Image -from torchvision.transforms import v2 - - -# ── Inlined from vitdet.py ───────────────────────────────────────────────── - -def init_t_xy(end_x, end_y, scale=1.0, offset=0): - t = torch.arange(end_x * end_y, dtype=torch.float32) - t_x = (t % end_x).float() - t_y = torch.div(t, end_x, rounding_mode="floor").float() - return t_x * scale + offset, t_y * scale + offset - - -def compute_axial_cis(dim, end_x, end_y, theta=10000.0, scale_pos=1.0, offset=0): - freqs_x = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim)) - freqs_y = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim)) - t_x, t_y = init_t_xy(end_x, end_y, scale_pos, offset) - freqs_x = torch.outer(t_x, freqs_x) - freqs_y = torch.outer(t_y, freqs_y) - freqs_cis_x = torch.polar(torch.ones_like(freqs_x), freqs_x) - freqs_cis_y = torch.polar(torch.ones_like(freqs_y), freqs_y) - return torch.cat([freqs_cis_x, freqs_cis_y], dim=-1) - - -def reshape_for_broadcast(freqs_cis, x): - ndim = x.ndim - shape = [d if i >= ndim - 2 else 1 for i, d in enumerate(x.shape)] - return freqs_cis.view(*shape) - - -def apply_rotary_enc(xq, xk, freqs_cis, repeat_freqs_k=False): - xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) - xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2)) if xk.shape[-2] != 0 else None - freqs_cis = reshape_for_broadcast(freqs_cis, xq_) - xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3) - if xk_ is None: - return xq_out.type_as(xq), xk - if repeat_freqs_k: - r = xk_.shape[-2] // xq_.shape[-2] - freqs_cis = freqs_cis.repeat(*([1] * (freqs_cis.ndim - 2)), r, 1) - xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3) - return xq_out.type_as(xq), xk_out.type_as(xk) - - -def window_partition(x, window_size): - B, H, W, C = x.shape - pad_h = (window_size - H % window_size) % window_size - pad_w = (window_size - W % window_size) % window_size - if pad_h > 0 or pad_w > 0: - x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h)) - Hp, Wp = H + pad_h, W + pad_w - x = x.view(B, Hp // window_size, window_size, Wp // window_size, window_size, C) - windows = x.permute(0, 1, 3, 2, 4, 5).reshape(-1, window_size, window_size, C) - return windows, (Hp, Wp) - - -def window_unpartition(windows, window_size, pad_hw, hw): - Hp, Wp = pad_hw - H, W = hw - B = windows.shape[0] // (Hp * Wp // window_size // window_size) - x = windows.reshape(B, Hp // window_size, Wp // window_size, window_size, window_size, -1) - x = x.permute(0, 1, 3, 2, 4, 5).reshape(B, Hp, Wp, -1) - if Hp > H or Wp > W: - x = x[:, :H, :W, :] - return x - - -def get_abs_pos(abs_pos, has_cls_token, hw, tiling=False): - h, w = hw - if has_cls_token: - abs_pos = abs_pos[:, 1:] - xy_num = abs_pos.shape[1] - size = int(math.sqrt(xy_num)) - if size != h or size != w: - new_abs_pos = abs_pos.reshape(1, size, size, -1).permute(0, 3, 1, 2) - if tiling: - new_abs_pos = new_abs_pos.tile( - [1, 1] + [x // y + 1 for x, y in zip((h, w), new_abs_pos.shape[2:])] - )[:, :, :h, :w] - else: - new_abs_pos = F.interpolate(new_abs_pos, size=(h, w), mode="bicubic", align_corners=False) - return new_abs_pos.permute(0, 2, 3, 1) - return abs_pos.reshape(1, h, w, -1) - - -# ── Minimal ViT block implementation ────────────────────────────────────── - -try: - from timm.layers import Mlp -except ImportError: - class Mlp(nn.Module): - def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features - self.fc1 = nn.Linear(in_features, hidden_features) - self.act = act_layer() - self.fc2 = nn.Linear(hidden_features, out_features) - def forward(self, x): - return self.fc2(self.act(self.fc1(x))) - - -def save_tensor(path, t): - t = t.detach().cpu().float().contiguous() - os.makedirs(os.path.dirname(path) or ".", exist_ok=True) - with open(path + ".bin", "wb") as f: - f.write(t.numpy().tobytes()) - with open(path + ".shape", "w") as f: - f.write(",".join(str(d) for d in t.shape)) - print(f" saved {path} shape={list(t.shape)}") - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--checkpoint", required=True) - parser.add_argument("--image", required=True) - parser.add_argument("--outdir", default="tests/ref") - args = parser.parse_args() - - os.makedirs(args.outdir, exist_ok=True) - device = "cpu" - - # ── Load checkpoint ───────────────────────────────────────────────── - print("Loading checkpoint...") - ckpt = torch.load(args.checkpoint, map_location=device, weights_only=False) - if "model" in ckpt: - ckpt = ckpt["model"] - - # Extract ViT weights - vit_prefix = "detector.backbone.vision_backbone.trunk." - vit_sd = {k[len(vit_prefix):]: v for k, v in ckpt.items() if k.startswith(vit_prefix)} - print(f" Found {len(vit_sd)} ViT keys") - - # ── Load and preprocess image ───────────────────────────────────── - print(f"Loading image: {args.image}") - img = Image.open(args.image).convert("RGB") - transform = v2.Compose([ - v2.ToDtype(torch.uint8, scale=True), - v2.Resize(size=(1008, 1008)), - v2.ToDtype(torch.float32, scale=True), - v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), - ]) - img_tensor = v2.functional.to_image(img) - img_preprocessed = transform(img_tensor).unsqueeze(0).to(device) - save_tensor(os.path.join(args.outdir, "preprocessed"), img_preprocessed) - - # ── Manual ViT forward pass ─────────────────────────────────────── - with torch.no_grad(): - # Config - E = 1024 - NH = 16 - HD = 64 - depth = 32 - MLP_DIM = 4736 - WS = 24 - global_blocks = {7, 15, 23, 31} - - # 1. Patch embedding: Conv2d(3, 1024, 14, 14, bias=False) - patch_w = vit_sd["patch_embed.proj.weight"] # [1024, 3, 14, 14] - x = F.conv2d(img_preprocessed, patch_w, stride=14) # [1, 1024, 72, 72] - x = x.permute(0, 2, 3, 1) # [1, 72, 72, 1024] - save_tensor(os.path.join(args.outdir, "patch_embed"), x) - - # 2. Positional embedding (tiled) - pos_embed = vit_sd["pos_embed"] # [1, 577, 1024] - pos = get_abs_pos(pos_embed, has_cls_token=True, hw=(72, 72), tiling=True) - save_tensor(os.path.join(args.outdir, "pos_embed_tiled"), pos) - x = x + pos - save_tensor(os.path.join(args.outdir, "after_pos_embed"), x) - - # 3. ln_pre - ln_pre_w = vit_sd["ln_pre.weight"] - ln_pre_b = vit_sd["ln_pre.bias"] - x = F.layer_norm(x, [E], ln_pre_w, ln_pre_b, eps=1e-5) - save_tensor(os.path.join(args.outdir, "after_ln_pre"), x) - - # 4. Transformer blocks - for blk_idx in range(depth): - prefix = f"blocks.{blk_idx}." - is_global = blk_idx in global_blocks - ws = 0 if is_global else WS - - # Pre-norm - n1_w = vit_sd[prefix + "norm1.weight"] - n1_b = vit_sd[prefix + "norm1.bias"] - shortcut = x - xn = F.layer_norm(x, [E], n1_w, n1_b, eps=1e-5) - - # Window partition - if ws > 0: - H, W = xn.shape[1], xn.shape[2] - xn, pad_hw = window_partition(xn, ws) - - # Attention - B_cur, Hc, Wc, _ = xn.shape - L = Hc * Wc - qkv_w = vit_sd[prefix + "attn.qkv.weight"] # [3*E, E] - qkv_b = vit_sd[prefix + "attn.qkv.bias"] # [3*E] - qkv = F.linear(xn, qkv_w, qkv_b).reshape(B_cur, L, 3, NH, HD) - q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0) # each [B, NH, L, HD] - - # RoPE - freqs_key = prefix + "attn.freqs_cis" - if freqs_key in vit_sd: - freqs = vit_sd[freqs_key] - q, k = apply_rotary_enc(q, k, freqs) - - # Scaled dot-product attention - attn_out = F.scaled_dot_product_attention(q, k, v) - # [B, NH, L, HD] → [B, H, W, E] - attn_out = attn_out.view(B_cur, NH, Hc, Wc, HD).permute(0, 2, 3, 1, 4).reshape(B_cur, Hc, Wc, E) - - # Output projection - proj_w = vit_sd[prefix + "attn.proj.weight"] - proj_b = vit_sd[prefix + "attn.proj.bias"] - attn_out = F.linear(attn_out, proj_w, proj_b) - - # Window unpartition - if ws > 0: - attn_out = window_unpartition(attn_out, ws, pad_hw, (H, W)) - - # Residual - x = shortcut + attn_out - - # FFN - n2_w = vit_sd[prefix + "norm2.weight"] - n2_b = vit_sd[prefix + "norm2.bias"] - shortcut = x - xn = F.layer_norm(x, [E], n2_w, n2_b, eps=1e-5) - - fc1_w = vit_sd[prefix + "mlp.fc1.weight"] - fc1_b = vit_sd[prefix + "mlp.fc1.bias"] - fc2_w = vit_sd[prefix + "mlp.fc2.weight"] - fc2_b = vit_sd[prefix + "mlp.fc2.bias"] - h = F.linear(xn, fc1_w, fc1_b) - h = F.gelu(h) - h = F.linear(h, fc2_w, fc2_b) - - x = shortcut + h - - if blk_idx == 0: - save_tensor(os.path.join(args.outdir, "block_0_out"), x) - if blk_idx == 7: - save_tensor(os.path.join(args.outdir, "block_7_out"), x) - - # Final output (x is [1, 72, 72, 1024]) - save_tensor(os.path.join(args.outdir, "vit_final"), x) - - # Permute to [B, C, H, W] - feats = x.permute(0, 3, 1, 2) # [1, 1024, 72, 72] - save_tensor(os.path.join(args.outdir, "vit_output_bchw"), feats) - - # ── Neck (SimpleFPN) ────────────────────────────────────────────── - print("\nDumping neck outputs...") - neck_prefix = "detector.backbone.vision_backbone.convs." - neck_sd = {k[len(neck_prefix):]: v for k, v in ckpt.items() if k.startswith(neck_prefix)} - - with torch.no_grad(): - # Scale 0 (4x): ConvTranspose(1024→512, k=2, s=2) → GELU → ConvTranspose(512→256, k=2, s=2) → Conv1x1 → Conv3x3 - s0 = F.conv_transpose2d(feats, neck_sd["0.dconv_2x2_0.weight"], neck_sd["0.dconv_2x2_0.bias"], stride=2) - s0 = F.gelu(s0) - s0 = F.conv_transpose2d(s0, neck_sd["0.dconv_2x2_1.weight"], neck_sd["0.dconv_2x2_1.bias"], stride=2) - s0 = F.conv2d(s0, neck_sd["0.conv_1x1.weight"], neck_sd["0.conv_1x1.bias"]) - s0 = F.conv2d(s0, neck_sd["0.conv_3x3.weight"], neck_sd["0.conv_3x3.bias"], padding=1) - save_tensor(os.path.join(args.outdir, "neck_det_0"), s0) - - # Scale 1 (2x) - s1 = F.conv_transpose2d(feats, neck_sd["1.dconv_2x2.weight"], neck_sd["1.dconv_2x2.bias"], stride=2) - s1 = F.conv2d(s1, neck_sd["1.conv_1x1.weight"], neck_sd["1.conv_1x1.bias"]) - s1 = F.conv2d(s1, neck_sd["1.conv_3x3.weight"], neck_sd["1.conv_3x3.bias"], padding=1) - save_tensor(os.path.join(args.outdir, "neck_det_1"), s1) - - # Scale 2 (1x) - s2 = F.conv2d(feats, neck_sd["2.conv_1x1.weight"], neck_sd["2.conv_1x1.bias"]) - s2 = F.conv2d(s2, neck_sd["2.conv_3x3.weight"], neck_sd["2.conv_3x3.bias"], padding=1) - save_tensor(os.path.join(args.outdir, "neck_det_2"), s2) - - # Scale 3 (0.5x) - s3 = F.max_pool2d(feats, kernel_size=2, stride=2) - s3 = F.conv2d(s3, neck_sd["3.conv_1x1.weight"], neck_sd["3.conv_1x1.bias"]) - s3 = F.conv2d(s3, neck_sd["3.conv_3x3.weight"], neck_sd["3.conv_3x3.bias"], padding=1) - save_tensor(os.path.join(args.outdir, "neck_det_3"), s3) - - print("\nDone!") - - -if __name__ == "__main__": - main() diff --git a/tests/llama.jpg b/tests/llama.jpg deleted file mode 100644 index 9d5116e..0000000 Binary files a/tests/llama.jpg and /dev/null differ diff --git a/tests/phase4_prompts.tsv b/tests/phase4_prompts.tsv deleted file mode 100644 index 3d8fae8..0000000 --- a/tests/phase4_prompts.tsv +++ /dev/null @@ -1,5 +0,0 @@ -bus yellow school bus -cat a photo of a cat -apostrophe don't stop -empty -long a person riding a bicycle on a sunny day in the park diff --git a/tests/phase6_cases.tsv b/tests/phase6_cases.tsv deleted file mode 100644 index d4abaef..0000000 --- a/tests/phase6_cases.tsv +++ /dev/null @@ -1,3 +0,0 @@ -point_single 0 504:504 -box_only 0 180:160:820:860 -box_and_points 1 620:640 260:260 180:160:820:860 diff --git a/tests/phase7_cases.tsv b/tests/phase7_cases.tsv deleted file mode 100644 index 16c7a47..0000000 --- a/tests/phase7_cases.tsv +++ /dev/null @@ -1,3 +0,0 @@ -single_mem_no_ptr Single memory, no pointer point_single 0 - -cond_prev_with_ptr Conditioning and previous memory with pointers point_single|box_only 0|1 1|2 -two_prev_with_ptr Two previous memories with pointers point_single|box_only 1|2 1|2 diff --git a/tests/test_debug_encoder.cpp b/tests/test_debug_encoder.cpp deleted file mode 100644 index bd42f94..0000000 --- a/tests/test_debug_encoder.cpp +++ /dev/null @@ -1,242 +0,0 @@ -/** - * Debug test for image encoder: loads Python-preprocessed image, - * runs through C++ ViT, dumps all intermediate tensors for comparison. - * - * Usage: - * ./test_debug_encoder - * - * ref_dir should contain preprocessed.bin/.shape from dump_phase3_reference.py. - * The test uses sam3_encode_image_from_preprocessed() to ensure identical input. - */ -#include "sam3.h" -#include "test_utils.h" - -#include -#include -#include -#include -#include - -// Compare C++ ggml tensor (dumped as [ne0, ne1, ne2]) vs Python NHWC [1, H, W, E] -// ggml stores: flat[e + w*E + h*E*W] where ne0=E, ne1=W, ne2=H -// Python NHWC: flat[h*W*E + w*E + e] (with optional batch dim) -// These are actually the same flat layout! -static compare_result compare_ggml_vs_nhwc(const ref_tensor_f32& cpp, - const ref_tensor_f32& ref) { - return compare_tensors(cpp.data.data(), ref.data.data(), - std::min(cpp.numel(), ref.numel())); -} - -// Compare C++ ggml tensor (dumped as [C, W, H]) vs Python NCHW [1, C, H, W] -// ggml stores: flat[c + w*C + h*C*W] -// Python NCHW: flat[c*H*W + h*W + w] -// These differ! Need transposition. -static compare_result compare_ggml_vs_nchw(const ref_tensor_f32& cpp, - const ref_tensor_f32& ref) { - if (cpp.shape.size() < 3 || ref.data.empty()) { - compare_result r; - r.max_diff = 999.0f; - return r; - } - int C = cpp.shape[0], W = cpp.shape[1], H = cpp.shape[2]; - - // Transpose ggml → NCHW - std::vector transposed(C * W * H); - for (int c = 0; c < C; ++c) - for (int h = 0; h < H; ++h) - for (int w = 0; w < W; ++w) - transposed[c * H * W + h * W + w] = cpp.data[c + w * C + h * C * W]; - - return compare_tensors(transposed.data(), ref.data.data(), - std::min((int)transposed.size(), ref.numel())); -} - -static void print_result(const char* name, const compare_result& r, float atol) { - const char* status = (r.max_diff < atol) ? "[PASS]" : "[FAIL]"; - fprintf(stderr, " %s %-40s max=%.6e mean=%.6e cos=%.8f\n", - status, name, r.max_diff, r.mean_diff, r.cosine_sim); - if (r.max_diff >= atol) { - fprintf(stderr, " worst_idx=%d a=%.6e b=%.6e\n", - r.worst_index, r.worst_a, r.worst_b); - } -} - -int main(int argc, char** argv) { - if (argc < 4) { - fprintf(stderr, "Usage: %s \n", argv[0]); - fprintf(stderr, "\nref_dir: directory with Python reference tensors\n"); - fprintf(stderr, "output_dir: where to dump C++ tensors\n"); - return 1; - } - - const std::string model_path = argv[1]; - const std::string ref_dir = argv[2]; - const std::string out_dir = argv[3]; - ensure_dir(out_dir); - - // ═══ Load model ═══ - fprintf(stderr, "\n═══ Loading model ═══\n"); - sam3_params params; - params.model_path = model_path; - params.use_gpu = false; // CPU for determinism - params.n_threads = 4; - - auto model = sam3_load_model(params); - if (!model) { - fprintf(stderr, "Failed to load model\n"); - return 1; - } - - auto state = sam3_create_state(*model, params); - if (!state) { - fprintf(stderr, "Failed to create state\n"); - return 1; - } - - // ═══ Stage 1: Load Python-preprocessed image ═══ - fprintf(stderr, "\n═══ Stage 1: Loading Python-preprocessed image ═══\n"); - auto ref_img = load_ref_f32(ref_dir + "/preprocessed"); - if (ref_img.data.empty()) { - fprintf(stderr, "Failed to load %s/preprocessed.bin\n", ref_dir.c_str()); - return 1; - } - fprintf(stderr, " Loaded preprocessed image: shape=["); - for (size_t i = 0; i < ref_img.shape.size(); ++i) { - if (i > 0) fprintf(stderr, ","); - fprintf(stderr, "%d", ref_img.shape[i]); - } - fprintf(stderr, "] (%d elements)\n", ref_img.numel()); - - const int img_size = 1008; - const int expected_elems = 3 * img_size * img_size; - if (ref_img.numel() != expected_elems) { - fprintf(stderr, " ERROR: expected %d elements, got %d\n", expected_elems, ref_img.numel()); - return 1; - } - - // ═══ Encode image using Python-preprocessed data ═══ - fprintf(stderr, "\n═══ Encoding image (from Python preprocessed data) ═══\n"); - bool ok = sam3_encode_image_from_preprocessed(*state, *model, ref_img.data.data(), img_size); - if (!ok) { - fprintf(stderr, "sam3_encode_image_from_preprocessed failed!\n"); - return 1; - } - - // ═══ Stage 2: Dump and compare ViT intermediate tensors ═══ - fprintf(stderr, "\n═══ Stage 2: Comparing ViT intermediate tensors ═══\n"); - - int n_pass = 0, n_fail = 0, n_skip = 0; - - // Helper to dump via sam3_dump_state_tensor and compare with ref in NHWC layout - auto check_nhwc = [&](const char* tensor_name, const char* ref_name, float atol) { - // Dump C++ tensor - sam3_dump_state_tensor(*state, tensor_name, out_dir + "/" + tensor_name); - - // Load both - auto cpp = load_ref_f32(out_dir + "/" + tensor_name); - auto ref = load_ref_f32(ref_dir + "/" + ref_name); - - if (cpp.data.empty()) { - fprintf(stderr, " [SKIP] %s — not found in C++ state\n", tensor_name); - n_skip++; - return; - } - if (ref.data.empty()) { - fprintf(stderr, " [SKIP] %s — no Python reference '%s'\n", tensor_name, ref_name); - n_skip++; - return; - } - - auto r = compare_ggml_vs_nhwc(cpp, ref); - print_result(ref_name, r, atol); - if (r.max_diff < atol) n_pass++; else n_fail++; - }; - - auto check_nchw = [&](const char* tensor_name, const char* ref_name, float atol) { - sam3_dump_state_tensor(*state, tensor_name, out_dir + "/" + tensor_name); - - auto cpp = load_ref_f32(out_dir + "/" + tensor_name); - auto ref = load_ref_f32(ref_dir + "/" + ref_name); - - if (cpp.data.empty()) { - fprintf(stderr, " [SKIP] %s — not found in C++ state\n", tensor_name); - n_skip++; - return; - } - if (ref.data.empty()) { - fprintf(stderr, " [SKIP] %s — no Python reference '%s'\n", tensor_name, ref_name); - n_skip++; - return; - } - - auto r = compare_ggml_vs_nchw(cpp, ref); - print_result(ref_name, r, atol); - if (r.max_diff < atol) n_pass++; else n_fail++; - }; - - // ── Patch Embedding ── - fprintf(stderr, "\n--- Patch Embedding ---\n"); - check_nhwc("dbg_patch_embed", "patch_embed", 1e-4f); - - // ── After Positional Embedding ── - fprintf(stderr, "\n--- After Positional Embedding ---\n"); - check_nhwc("dbg_after_pos_embed", "after_pos_embed", 1e-4f); - - // ── After LayerNorm Pre ── - fprintf(stderr, "\n--- After LayerNorm Pre ---\n"); - check_nhwc("dbg_after_ln_pre", "after_ln_pre", 1e-4f); - - // ── Per-block outputs ── - // Tolerance is progressive: early blocks have tiny error, later blocks accumulate - // f32 drift on values of magnitude ~100-270 across 32 blocks. - // Max absolute error grows but relative error stays < 0.1% and cosine ≈ 1.0. - fprintf(stderr, "\n--- ViT Block Outputs ---\n"); - for (int i = 0; i < 32; ++i) { - char cpp_name[64], ref_name[64]; - snprintf(cpp_name, sizeof(cpp_name), "dbg_block_%d_out", i); - snprintf(ref_name, sizeof(ref_name), "block_%d_out", i); - - // Check if Python reference exists for this block - auto ref = load_ref_f32(ref_dir + "/" + ref_name); - if (ref.data.empty()) continue; - - // Realistic tolerance for f32 accumulation over transformer blocks. - // Values reach magnitude ~270 by block 20, so 0.05 absolute tolerance - // corresponds to ~0.02% relative error. - float atol = 1e-3f + i * 2e-3f; - check_nhwc(cpp_name, ref_name, atol); - } - - // ── Final ViT output (NCHW comparison) ── - fprintf(stderr, "\n--- Final ViT Output ---\n"); - check_nchw("vit_output", "vit_output_bchw", 0.05f); - - // ── Neck outputs (NCHW comparison) ── - fprintf(stderr, "\n--- Neck Outputs (detector) ---\n"); - for (int i = 0; i < 4; ++i) { - char name[32]; - snprintf(name, sizeof(name), "neck_det_%d", i); - check_nchw(name, name, 5e-3f); - } - - fprintf(stderr, "\n--- Neck Outputs (tracker) ---\n"); - for (int i = 0; i < 4; ++i) { - char name[32]; - snprintf(name, sizeof(name), "neck_trk_%d", i); - // Check if ref exists for tracker path - auto ref = load_ref_f32(ref_dir + "/" + name); - if (ref.data.empty()) continue; - check_nchw(name, name, 5e-3f); - } - - // ═══ Summary ═══ - fprintf(stderr, "\n═══════════════════════════════════════════\n"); - fprintf(stderr, "Results: %d passed, %d failed, %d skipped\n", n_pass, n_fail, n_skip); - fprintf(stderr, "═══════════════════════════════════════════\n"); - - state.reset(); - sam3_free_model(*model); - model.reset(); - - return n_fail > 0 ? 1 : 0; -} diff --git a/tests/test_debug_fenc.cpp b/tests/test_debug_fenc.cpp deleted file mode 100644 index 2dd3343..0000000 --- a/tests/test_debug_fenc.cpp +++ /dev/null @@ -1,174 +0,0 @@ -#include "sam3.h" - -#include "test_utils.h" - -#include -#include -#include - -struct tensor_case { - std::string name; - std::string label; - float atol; -}; - -static std::vector build_cases() { - std::vector cases; - for (int i = 0; i < 6; ++i) { - char name[64], label[64]; - snprintf(name, sizeof(name), "fenc_layer%d_out", i); - snprintf(label, sizeof(label), "Fusion encoder layer %d", i); - // Tolerance: early layers should be tight, later layers allow more accumulation - float atol = (i < 3) ? 1e-4f : 2e-4f; - cases.push_back({name, label, atol}); - } - return cases; -} - -int main(int argc, char ** argv) { - if (argc < 3) { - fprintf(stderr, - "Usage: %s [n_threads]\n" - "\n" - " ref_dir : directory with Python-dumped fenc tensors\n" - " (from dump_fenc_from_package.py)\n" - " model_path : path to sam3_f32.ggml weights\n" - " n_threads : optional (default 1)\n", - argv[0]); - return 1; - } - - const std::string ref_dir = argv[1]; - const std::string model_path = argv[2]; - const int n_threads = (argc > 3) ? std::atoi(argv[3]) : 1; - - const std::string cpp_dir = ref_dir + "/cpp_out_fenc"; - if (!ensure_dir(cpp_dir)) { - fprintf(stderr, "Failed to create output dir %s\n", cpp_dir.c_str()); - return 1; - } - - // ── Load model ── - sam3_params params; - params.model_path = model_path; - params.n_threads = n_threads; - params.use_gpu = false; - - fprintf(stderr, "Loading model from %s...\n", model_path.c_str()); - auto model = sam3_load_model(params); - if (!model) { - fprintf(stderr, "Failed to load model\n"); - return 1; - } - - // ── Run fenc-only test ── - fprintf(stderr, "\nRunning fusion encoder with Python-dumped inputs from %s\n", ref_dir.c_str()); - if (!sam3_test_fenc_only(*model, ref_dir, cpp_dir, n_threads)) { - fprintf(stderr, "sam3_test_fenc_only failed\n"); - return 1; - } - - // ── Compare per-layer outputs ── - fprintf(stderr, "\n══════════════════════════════════════════════════════════════\n"); - fprintf(stderr, " Fusion Encoder: Python vs C++ Comparison\n"); - fprintf(stderr, "══════════════════════════════════════════════════════════════\n\n"); - - // First verify inputs are loaded correctly - { - auto ref_img = load_ref_f32(ref_dir + "/fenc_input_tgt"); - auto cpp_img = load_ref_f32(cpp_dir + "/fenc_img_input"); - if (!ref_img.data.empty() && !cpp_img.data.empty() && - ref_img.numel() == cpp_img.numel()) { - auto r = compare_tensors(cpp_img.data.data(), ref_img.data.data(), - ref_img.numel(), 1e-6f); - fprintf(stderr, " Input check: img_feats max_diff=%.2e (should be ~0)\n", - r.max_diff); - } - - auto ref_pos = load_ref_f32(ref_dir + "/fenc_input_pos"); - auto cpp_pos = load_ref_f32(cpp_dir + "/fenc_pos_input"); - if (!ref_pos.data.empty() && !cpp_pos.data.empty() && - ref_pos.numel() == cpp_pos.numel()) { - auto r = compare_tensors(cpp_pos.data.data(), ref_pos.data.data(), - ref_pos.numel(), 1e-6f); - fprintf(stderr, " Input check: pos_embed max_diff=%.2e (should be ~0)\n", - r.max_diff); - } - - auto ref_prompt = load_ref_f32(ref_dir + "/fenc_input_prompt"); - auto cpp_prompt = load_ref_f32(cpp_dir + "/fenc_prompt_input"); - if (!ref_prompt.data.empty() && !cpp_prompt.data.empty() && - ref_prompt.numel() == cpp_prompt.numel()) { - auto r = compare_tensors(cpp_prompt.data.data(), ref_prompt.data.data(), - ref_prompt.numel(), 1e-6f); - fprintf(stderr, " Input check: prompt max_diff=%.2e (should be ~0)\n", - r.max_diff); - } - } - - fprintf(stderr, "\n"); - - // Compare per-layer outputs - const auto cases = build_cases(); - bool overall_ok = true; - int n_pass = 0; - - fprintf(stderr, "| %-30s | %10s | %10s | %10s | %10s | %6s |\n", - "Layer", "max_diff", "mean_diff", "cosine", "tolerance", "status"); - fprintf(stderr, "|%s|%s|%s|%s|%s|%s|\n", - "--------------------------------", "------------", "------------", - "------------", "------------", "--------"); - - for (const auto & tc : cases) { - auto ref = load_ref_f32(ref_dir + "/" + tc.name); - auto cpp = load_ref_f32(cpp_dir + "/" + tc.name); - - if (ref.data.empty()) { - fprintf(stderr, "| %-30s | %10s | %10s | %10s | %10.1e | %-6s |\n", - tc.label.c_str(), "N/A", "N/A", "N/A", tc.atol, "SKIP"); - fprintf(stderr, " (missing Python reference %s/%s)\n", ref_dir.c_str(), tc.name.c_str()); - continue; - } - if (cpp.data.empty()) { - fprintf(stderr, "| %-30s | %10s | %10s | %10s | %10.1e | %-6s |\n", - tc.label.c_str(), "N/A", "N/A", "N/A", tc.atol, "SKIP"); - fprintf(stderr, " (missing C++ output %s/%s)\n", cpp_dir.c_str(), tc.name.c_str()); - continue; - } - - if (ref.numel() != cpp.numel()) { - fprintf(stderr, "| %-30s | %10s | %10s | %10s | %10.1e | %-6s |\n", - tc.label.c_str(), "N/A", "N/A", "N/A", tc.atol, "FAIL"); - fprintf(stderr, " SHAPE MISMATCH: ref=%d cpp=%d\n", ref.numel(), cpp.numel()); - overall_ok = false; - continue; - } - - auto r = compare_tensors(cpp.data.data(), ref.data.data(), ref.numel(), tc.atol); - bool pass = r.max_diff <= tc.atol; - if (pass) n_pass++; - if (!pass) overall_ok = false; - - fprintf(stderr, "| %-30s | %10.3e | %10.3e | %10.8f | %10.1e | %-6s |\n", - tc.label.c_str(), - r.max_diff, r.mean_diff, r.cosine_sim, - tc.atol, pass ? "PASS" : "FAIL"); - - if (!pass) { - fprintf(stderr, " worst idx=%d ref=%.6g cpp=%.6g bad=%d/%d\n", - r.worst_index, r.worst_b, r.worst_a, r.n_bad, r.n_total); - } - } - - fprintf(stderr, "\n%d/%zu layers passed\n", n_pass, cases.size()); - fprintf(stderr, "Overall: %s\n", overall_ok ? "PASS" : "FAIL"); - - if (!overall_ok) { - fprintf(stderr, "\n>>> DIVERGENCE DETECTED. Next steps:\n"); - fprintf(stderr, " 1. If layer 0 fails: bug is in self-attention or cross-attention\n"); - fprintf(stderr, " 2. Add finer-grained dumps (after SA, after CA, after FFN)\n"); - fprintf(stderr, " 3. Check weight loading in convert_sam3_to_ggml.py\n"); - } - - return overall_ok ? 0 : 1; -} diff --git a/tests/test_e2e_pcs.cpp b/tests/test_e2e_pcs.cpp deleted file mode 100644 index 3b74e89..0000000 --- a/tests/test_e2e_pcs.cpp +++ /dev/null @@ -1,468 +0,0 @@ -/** - * End-to-end PCS pipeline test: loads Python-preprocessed image, - * runs through C++ ViT + neck + text encoder + fusion encoder + DETR decoder - * + segmentation head, dumps all intermediate tensors for comparison against - * Python reference. - * - * Usage: - * ./test_e2e_pcs - * - * ref_dir should contain pcs_ref/ with tensors from dump_e2e_pcs.py: - * - preprocessed_chw.bin/.shape - * - token_ids.bin/.shape - * - text_features.bin/.shape - * - neck_det_*.bin/.shape - * - fenc_layer*_out.bin/.shape - * - ddec_layer*_out.bin/.shape - * - pred_logits.bin/.shape, pred_boxes.bin/.shape - * - seg_pred_masks.bin/.shape - * etc. - */ -#include "sam3.h" -#include "test_utils.h" - -#include -#include -#include -#include -#include - -// ── Metric row (same as test_e2e_pvs.cpp) ── - -struct metric_row { - std::string stage; - std::string name; - std::string shape_str; - float mae = 0.0f; - float max_abs = 0.0f; - float mean_rel = 0.0f; - float cosine = 0.0f; - float p95 = 0.0f; - float p99 = 0.0f; - bool ok = false; - float tolerance = 0.0f; -}; - -static metric_row compute_full_metrics(const float * a, const float * b, int n, float tol) { - metric_row m; - m.tolerance = tol; - if (n == 0) return m; - - double sum_diff = 0.0, sum_rel = 0.0; - double dot_ab = 0.0, dot_aa = 0.0, dot_bb = 0.0; - float max_d = 0.0f; - - std::vector diffs(n); - for (int i = 0; i < n; ++i) { - float d = fabsf(a[i] - b[i]); - diffs[i] = d; - sum_diff += d; - if (d > max_d) max_d = d; - float denom = fabsf(b[i]) + 1e-8f; - sum_rel += d / denom; - dot_ab += (double)a[i] * (double)b[i]; - dot_aa += (double)a[i] * (double)a[i]; - dot_bb += (double)b[i] * (double)b[i]; - } - - m.mae = (float)(sum_diff / n); - m.max_abs = max_d; - m.mean_rel = (float)(sum_rel / n); - double ddenom = sqrt(dot_aa) * sqrt(dot_bb); - m.cosine = ddenom > 0.0 ? (float)(dot_ab / ddenom) : 0.0f; - - std::sort(diffs.begin(), diffs.end()); - m.p95 = diffs[(int)(0.95 * n)]; - m.p99 = diffs[(int)(0.99 * n)]; - m.ok = max_d <= tol; - return m; -} - -// Compare raw flat data (same byte layout between Python and C++) -static metric_row compare_flat(const char * stage, const char * label, - const std::string & cpp_path, - const std::string & ref_path, - float atol) { - auto cpp = load_ref_f32(cpp_path); - auto ref = load_ref_f32(ref_path); - metric_row m; - m.stage = stage; - m.name = label; - m.tolerance = atol; - if (cpp.data.empty() || ref.data.empty()) { - fprintf(stderr, " [SKIP] %s (cpp=%s ref=%s)\n", label, - cpp.data.empty() ? "missing" : "ok", - ref.data.empty() ? "missing" : "ok"); - return m; - } - std::string s; - for (size_t i = 0; i < ref.shape.size(); ++i) { - if (i > 0) s += ","; - s += std::to_string(ref.shape[i]); - } - m.shape_str = s; - int n = std::min(cpp.numel(), ref.numel()); - return compute_full_metrics(cpp.data.data(), ref.data.data(), n, atol); -} - -int main(int argc, char ** argv) { - if (argc < 3) { - fprintf(stderr, "Usage: %s \n", argv[0]); - fprintf(stderr, "\nref_dir should contain pcs_ref/ from dump_e2e_pcs.py\n"); - return 1; - } - - const std::string model_path = argv[1]; - const std::string ref_dir = std::string(argv[2]) + "/pcs_ref"; - const std::string cpp_out = std::string(argv[2]) + "/pcs_cpp"; - ensure_dir(cpp_out); - - // ═══ Load model ═══ - fprintf(stderr, "\n═══ Loading model ═══\n"); - sam3_params params; - params.model_path = model_path; - params.use_gpu = false; - params.n_threads = 4; - - auto model = sam3_load_model(params); - if (!model) { - fprintf(stderr, "Failed to load model\n"); - return 1; - } - - auto state = sam3_create_state(*model, params); - if (!state) { - fprintf(stderr, "Failed to create state\n"); - return 1; - } - - // ═══ Stage 1: Load Python-preprocessed image ═══ - fprintf(stderr, "\n═══ Stage 1: Loading Python-preprocessed image ═══\n"); - auto ref_img = load_ref_f32(ref_dir + "/preprocessed_chw"); - if (ref_img.data.empty()) { - fprintf(stderr, "Failed to load %s/preprocessed_chw.bin\n", ref_dir.c_str()); - return 1; - } - fprintf(stderr, " Loaded preprocessed image: %d elements\n", ref_img.numel()); - - const int img_size = 1008; - bool ok = sam3_encode_image_from_preprocessed(*state, *model, ref_img.data.data(), img_size); - if (!ok) { - fprintf(stderr, "sam3_encode_image_from_preprocessed failed!\n"); - return 1; - } - fprintf(stderr, " Image encoded successfully\n"); - - // ═══ Stage 2: Dump ViT block outputs (from state debug tensors) ═══ - fprintf(stderr, "\n═══ Stage 2: ViT Backbone ═══\n"); - std::vector report; - - // ViT block outputs are stored in state during encoding - for (int i : {0, 7, 15, 23, 31}) { - char cpp_name[64], ref_name[64]; - snprintf(cpp_name, sizeof(cpp_name), "dbg_block_%d_out", i); - snprintf(ref_name, sizeof(ref_name), "vit_block_%02d_out", i); - - sam3_dump_state_tensor(*state, cpp_name, cpp_out + "/" + cpp_name); - auto cpp = load_ref_f32(cpp_out + "/" + cpp_name); - auto ref = load_ref_f32(ref_dir + "/" + ref_name); - if (cpp.data.empty() || ref.data.empty()) { - fprintf(stderr, " [SKIP] %s\n", ref_name); - continue; - } - // f32 accumulation through 32 blocks: allow up to 0.07 max error - int n = std::min(cpp.numel(), ref.numel()); - auto m = compute_full_metrics(cpp.data.data(), ref.data.data(), n, 0.07f); - m.stage = "ViT"; - m.name = ref_name; - char shape_s[64]; - snprintf(shape_s, sizeof(shape_s), "%d,%d,%d,%d", - ref.shape.size()>0?ref.shape[0]:0, ref.shape.size()>1?ref.shape[1]:0, - ref.shape.size()>2?ref.shape[2]:0, ref.shape.size()>3?ref.shape[3]:0); - m.shape_str = shape_s; - report.push_back(m); - } - - // ═══ Stage 3: Neck ═══ - fprintf(stderr, "\n═══ Stage 3: Neck (SimpleFPN) ═══\n"); - // Dump neck detector features — these are stored in state - for (int i = 0; i < 3; ++i) { - char tensor_name[32], ref_name[32]; - snprintf(tensor_name, sizeof(tensor_name), "neck_det_%d", i); - snprintf(ref_name, sizeof(ref_name), "neck_det_%d", i); - - sam3_dump_state_tensor(*state, tensor_name, cpp_out + "/" + tensor_name); - auto cpp_t = load_ref_f32(cpp_out + "/" + tensor_name); - auto ref_t = load_ref_f32(ref_dir + "/" + ref_name); - if (cpp_t.data.empty() || ref_t.data.empty()) { - fprintf(stderr, " [SKIP] %s\n", ref_name); - continue; - } - - // ggml stores [C, W, H, B], Python dumps NCHW as [C, W, H, B] via ggml layout - // Both should match in flat layout now - int n = std::min(cpp_t.numel(), ref_t.numel()); - auto m = compute_full_metrics(cpp_t.data.data(), ref_t.data.data(), n, 5e-3f); - m.stage = "Neck"; - m.name = ref_name; - std::string s; - for (size_t j = 0; j < ref_t.shape.size(); ++j) { - if (j > 0) s += ","; - s += std::to_string(ref_t.shape[j]); - } - m.shape_str = s; - report.push_back(m); - } - - // ═══ Stage 4-9: Text + Fusion + DETR + Segmentation (via sam3_test_dump_phase5) ═══ - fprintf(stderr, "\n═══ Stage 4-9: Text Encoder + Fusion + DETR + Segmentation ═══\n"); - - // Load reference token IDs (saved as float32 by Python dump) - auto ref_tokens_f = load_ref_f32(ref_dir + "/token_ids"); - if (ref_tokens_f.data.empty()) { - fprintf(stderr, "Failed to load token IDs from ref_dir\n"); - return 1; - } - // Convert float→int32 - std::vector token_ids(ref_tokens_f.numel()); - for (int i = 0; i < ref_tokens_f.numel(); ++i) { - token_ids[i] = (int32_t)ref_tokens_f.data[i]; - } - // Pad to 32 if needed - while ((int)token_ids.size() < 32) token_ids.push_back(0); - fprintf(stderr, " Token IDs: ["); - for (int i = 0; i < 32; ++i) { - if (i > 0) fprintf(stderr, ", "); - fprintf(stderr, "%d", token_ids[i]); - } - fprintf(stderr, "]\n"); - - // Run the full phase 5 dump (text encoder + fusion encoder + DETR decoder + seg head) - const std::string phase5_out = cpp_out + "/phase5"; - ensure_dir(phase5_out); - if (!sam3_test_dump_phase5(*model, *state, token_ids, phase5_out, params.n_threads)) { - fprintf(stderr, " [FAIL] sam3_test_dump_phase5 failed\n"); - return 1; - } - fprintf(stderr, " Phase 5 tensors dumped to %s\n", phase5_out.c_str()); - - // ═══ Compare text encoder ═══ - fprintf(stderr, "\n --- Text Encoder ---\n"); - { - auto m = compare_flat("TextEnc", "text_features", - phase5_out + "/text_features", - ref_dir + "/text_features", 1e-4f); - m.stage = "TextEnc"; - report.push_back(m); - } - - // ═══ Compare fusion encoder layers ═══ - fprintf(stderr, "\n --- Fusion Encoder ---\n"); - { - // Compare fenc_output (the final encoder output) - // C++ dumps as [D, 5184, 1], Python dumps as [5184, 1, 256] - // Both are flat [5184*256] floats in memory - auto m = compare_flat("FEnc", "fenc_output", - phase5_out + "/fenc_output", - ref_dir + "/fenc_output", 5e-3f); - m.stage = "FEnc"; - report.push_back(m); - } - for (int i = 0; i < 6; ++i) { - char name[64]; - snprintf(name, sizeof(name), "fenc_layer%d_out", i); - auto m = compare_flat("FEnc", name, - phase5_out + "/" + name, - ref_dir + "/" + name, 5e-3f); - m.stage = "FEnc"; - report.push_back(m); - } - - // ═══ Compare DETR decoder layers ═══ - fprintf(stderr, "\n --- DETR Decoder ---\n"); - for (int i = 0; i < 6; ++i) { - char name[64]; - snprintf(name, sizeof(name), "ddec_layer%d_out", i); - auto m = compare_flat("DDec", name, - phase5_out + "/" + name, - ref_dir + "/" + name, 0.1f); - m.stage = "DDec"; - report.push_back(m); - } - - // ═══ Compare scoring ═══ - fprintf(stderr, "\n --- Scoring ---\n"); - { - auto m = compare_flat("Score", "scoring_class_scores", - phase5_out + "/scoring_class_scores", - ref_dir + "/scoring_output", 0.5f); - m.stage = "Score"; - report.push_back(m); - } - - // ═══ Compare segmentation masks ═══ - fprintf(stderr, "\n --- Segmentation ---\n"); - { - auto m = compare_flat("SegHead", "seg_mask_logits", - phase5_out + "/seg_mask_logits", - ref_dir + "/seg_pred_masks", 1.0f); - m.stage = "SegHead"; - report.push_back(m); - } - - // ═══ Compare final model outputs ═══ - fprintf(stderr, "\n --- Final Outputs ---\n"); - { - // pred_boxes - auto m = compare_flat("Output", "pred_boxes", - phase5_out + "/ddec_pred_boxes", - ref_dir + "/pred_boxes", 0.1f); - m.stage = "Output"; - report.push_back(m); - } - { - // presence score - auto m = compare_flat("Output", "presence_logit", - phase5_out + "/ddec_presence_logit", - ref_dir + "/presence_logit", 0.5f); - m.stage = "Output"; - report.push_back(m); - } - - // ═══ Now run the actual sam3_segment_pcs and compare final results ═══ - fprintf(stderr, "\n═══ Running sam3_segment_pcs (end-to-end) ═══\n"); - - // Re-encode image (sam3_test_dump_phase5 may have invalidated state) - ok = sam3_encode_image_from_preprocessed(*state, *model, ref_img.data.data(), img_size); - if (!ok) { - fprintf(stderr, "Failed to re-encode image\n"); - return 1; - } - - sam3_pcs_params pcs; - pcs.text_prompt = "cat"; - pcs.score_threshold = 0.3f; - pcs.nms_threshold = 0.1f; - - sam3_result result = sam3_segment_pcs(*state, *model, pcs); - - fprintf(stderr, "\n C++ detections: %zu\n", result.detections.size()); - for (size_t i = 0; i < result.detections.size(); ++i) { - const auto & det = result.detections[i]; - fprintf(stderr, " Det %zu: score=%.4f box=[%.1f, %.1f, %.1f, %.1f]\n", - i, det.score, det.box.x0, det.box.y0, det.box.x1, det.box.y1); - } - - // Compare with Python results - auto py_scores = load_ref_f32(ref_dir + "/final_scores"); - auto py_boxes = load_ref_f32(ref_dir + "/final_boxes_xyxy"); - - if (!py_scores.data.empty() && !result.detections.empty()) { - fprintf(stderr, "\n Python detections: %d\n", py_scores.numel()); - for (int i = 0; i < py_scores.numel(); ++i) { - fprintf(stderr, " Py Det %d: score=%.4f box=[%.1f, %.1f, %.1f, %.1f]\n", - i, py_scores.data[i], - py_boxes.data[i*4+0], py_boxes.data[i*4+1], - py_boxes.data[i*4+2], py_boxes.data[i*4+3]); - } - - // Compare top detection - float score_diff = fabsf(result.detections[0].score - py_scores.data[0]); - float box_diff = fabsf(result.detections[0].box.x0 - py_boxes.data[0]) - + fabsf(result.detections[0].box.y0 - py_boxes.data[1]) - + fabsf(result.detections[0].box.x1 - py_boxes.data[2]) - + fabsf(result.detections[0].box.y1 - py_boxes.data[3]); - fprintf(stderr, "\n Score diff: %.4f\n", score_diff); - fprintf(stderr, " Box L1 diff: %.2f pixels\n", box_diff); - - metric_row m_score; - m_score.stage = "E2E"; - m_score.name = "score_match"; - m_score.max_abs = score_diff; - m_score.ok = score_diff < 0.05f; - m_score.tolerance = 0.05f; - report.push_back(m_score); - - metric_row m_box; - m_box.stage = "E2E"; - m_box.name = "box_match (L1 px)"; - m_box.max_abs = box_diff; - m_box.ok = box_diff < 50.0f; - m_box.tolerance = 50.0f; - report.push_back(m_box); - } - - // ═══ Compare masks ═══ - if (!result.detections.empty() && !result.detections[0].mask.data.empty()) { - auto py_mask = load_ref_f32(ref_dir + "/final_masks_binary"); - if (!py_mask.data.empty()) { - int mask_w = result.detections[0].mask.width; - int mask_h = result.detections[0].mask.height; - int n_pix = mask_w * mask_h; - int n_match = 0; - int n_total = std::min(n_pix, py_mask.numel()); - for (int i = 0; i < n_total; ++i) { - float cpp_val = result.detections[0].mask.data[i] > 0 ? 1.0f : 0.0f; - float py_val = py_mask.data[i]; - if (fabsf(cpp_val - py_val) < 0.5f) n_match++; - } - float iou_approx = (float)n_match / (float)n_total; - fprintf(stderr, "\n Mask agreement: %d/%d pixels (%.2f%%)\n", - n_match, n_total, 100.0f * iou_approx); - - metric_row m_mask; - m_mask.stage = "E2E"; - m_mask.name = "mask_agreement"; - m_mask.max_abs = 1.0f - iou_approx; - m_mask.ok = iou_approx > 0.95f; - m_mask.tolerance = 0.05f; - m_mask.cosine = iou_approx; - report.push_back(m_mask); - } - } - - // ═══ Print Full Report ═══ - fprintf(stderr, "\n"); - fprintf(stderr, "══════════════════════════════════════════════════════════════════════════════════════════════════════════════════════════\n"); - fprintf(stderr, " END-TO-END PCS PIPELINE COMPARISON: Python vs C++ (cat.jpg with text=\"cat\")\n"); - fprintf(stderr, "══════════════════════════════════════════════════════════════════════════════════════════════════════════════════════════\n\n"); - - int n_pass = 0, n_fail = 0, n_skip = 0; - fprintf(stderr, " %-8s %-8s %-30s %-20s %12s %12s %12s %12s %12s %12s\n", - "Status", "Stage", "Tensor", "Shape", "MAE", "Max", "RelErr", "Cosine", "P95", "P99"); - fprintf(stderr, " %-8s %-8s %-30s %-20s %12s %12s %12s %12s %12s %12s\n", - "------", "-----", "------", "-----", "---", "---", "------", "------", "---", "---"); - - for (const auto & m : report) { - if (m.shape_str.empty() && m.mae == 0.0f && m.max_abs == 0.0f && m.cosine == 0.0f) { - fprintf(stderr, " %-8s %-8s %-30s\n", "SKIP", m.stage.c_str(), m.name.c_str()); - n_skip++; - continue; - } - const char * status = m.ok ? "PASS" : "FAIL"; - fprintf(stderr, " %-8s %-8s %-30s %-20s %12.4e %12.4e %12.4e %12.8f %12.4e %12.4e\n", - status, m.stage.c_str(), m.name.c_str(), m.shape_str.c_str(), - m.mae, m.max_abs, m.mean_rel, m.cosine, m.p95, m.p99); - if (m.ok) n_pass++; - else n_fail++; - } - - fprintf(stderr, "\n══════════════════════════════════════════════════════════════════════════════════════════════════════════════════════════\n"); - fprintf(stderr, " TOTAL: %d PASS, %d FAIL, %d SKIP out of %d tensors\n", - n_pass, n_fail, n_skip, (int)report.size()); - fprintf(stderr, "══════════════════════════════════════════════════════════════════════════════════════════════════════════════════════════\n\n"); - - // Save mask for visual comparison - if (!result.detections.empty()) { - std::string mask_path = cpp_out + "/cat_mask.png"; - if (sam3_save_mask(result.detections[0].mask, mask_path)) { - fprintf(stderr, " Saved C++ mask: %s\n", mask_path.c_str()); - } - } - - state.reset(); - sam3_free_model(*model); - model.reset(); - - return n_fail > 0 ? 1 : 0; -} diff --git a/tests/test_e2e_pvs.cpp b/tests/test_e2e_pvs.cpp deleted file mode 100644 index 848cffa..0000000 --- a/tests/test_e2e_pvs.cpp +++ /dev/null @@ -1,333 +0,0 @@ -/** - * End-to-end PVS pipeline test: loads Python-preprocessed image, - * runs through C++ ViT + neck + SAM prompt encoder + mask decoder, - * dumps all intermediate tensors for comparison against Python reference. - * - * Usage: - * ./test_e2e_pvs - * - * ref_dir should contain: - * - preprocessed.bin/.shape (from dump_phase3_reference.py) - * - phase6/cat_box/ (from dump_phase6_reference.py) - */ -#include "sam3.h" -#include "test_utils.h" - -#include -#include -#include -#include -#include - -// Compare ggml [ne0, ne1, ne2] vs Python NHWC [1, H, W, E] — same flat layout -static compare_result compare_ggml_vs_nhwc(const ref_tensor_f32 & cpp, - const ref_tensor_f32 & ref) { - return compare_tensors(cpp.data.data(), ref.data.data(), - std::min(cpp.numel(), ref.numel())); -} - -// Compare ggml [C, W, H] vs Python NCHW [1, C, H, W] — need transposition -static compare_result compare_ggml_vs_nchw(const ref_tensor_f32 & cpp, - const ref_tensor_f32 & ref) { - if (cpp.shape.size() < 3 || ref.data.empty()) { - compare_result r; - r.max_diff = 999.0f; - return r; - } - int C = cpp.shape[0], W = cpp.shape[1], H = cpp.shape[2]; - - std::vector transposed(C * W * H); - for (int c = 0; c < C; ++c) - for (int h = 0; h < H; ++h) - for (int w = 0; w < W; ++w) - transposed[c * H * W + h * W + w] = cpp.data[c + w * C + h * C * W]; - - return compare_tensors(transposed.data(), ref.data.data(), - std::min((int) transposed.size(), ref.numel())); -} - -struct metric_row { - std::string stage; - std::string name; - std::string py_shape_str; - float mae = 0.0f; - float max_abs = 0.0f; - float mean_rel = 0.0f; - float cosine = 0.0f; - float p95 = 0.0f; - float p99 = 0.0f; - bool ok = false; - float tolerance = 0.0f; -}; - -static metric_row compute_full_metrics(const float * a, const float * b, int n, float tol) { - metric_row m; - m.tolerance = tol; - if (n == 0) return m; - - double sum_diff = 0.0; - double sum_rel = 0.0; - double dot_ab = 0.0, dot_aa = 0.0, dot_bb = 0.0; - float max_d = 0.0f; - - std::vector diffs(n); - for (int i = 0; i < n; ++i) { - float d = fabsf(a[i] - b[i]); - diffs[i] = d; - sum_diff += d; - if (d > max_d) max_d = d; - float denom = fabsf(b[i]) + 1e-8f; - sum_rel += d / denom; - dot_ab += (double)a[i] * (double)b[i]; - dot_aa += (double)a[i] * (double)a[i]; - dot_bb += (double)b[i] * (double)b[i]; - } - - m.mae = (float)(sum_diff / n); - m.max_abs = max_d; - m.mean_rel = (float)(sum_rel / n); - double denom = sqrt(dot_aa) * sqrt(dot_bb); - m.cosine = denom > 0.0 ? (float)(dot_ab / denom) : 0.0f; - - std::sort(diffs.begin(), diffs.end()); - m.p95 = diffs[(int)(0.95 * n)]; - m.p99 = diffs[(int)(0.99 * n)]; - m.ok = max_d <= tol; - return m; -} - -int main(int argc, char ** argv) { - if (argc < 3) { - fprintf(stderr, "Usage: %s \n", argv[0]); - fprintf(stderr, "\nref_dir: directory with Python reference tensors\n"); - fprintf(stderr, " Should contain preprocessed.bin/.shape and phase6/cat_box/\n"); - return 1; - } - - const std::string model_path = argv[1]; - const std::string ref_dir = argv[2]; - const std::string cpp_out = ref_dir + "/cpp_out"; - const std::string cpp_p6_out = ref_dir + "/cpp_out_phase6"; - ensure_dir(cpp_out); - ensure_dir(cpp_p6_out); - - // ═══ Load model ═══ - fprintf(stderr, "\n═══ Loading model ═══\n"); - sam3_params params; - params.model_path = model_path; - params.use_gpu = false; - params.n_threads = 4; - - auto model = sam3_load_model(params); - if (!model) { - fprintf(stderr, "Failed to load model\n"); - return 1; - } - - auto state = sam3_create_state(*model, params); - if (!state) { - fprintf(stderr, "Failed to create state\n"); - return 1; - } - - // ═══ Stage 1: Load Python-preprocessed image ═══ - fprintf(stderr, "\n═══ Stage 1: Loading Python-preprocessed image ═══\n"); - auto ref_img = load_ref_f32(ref_dir + "/preprocessed"); - if (ref_img.data.empty()) { - fprintf(stderr, "Failed to load %s/preprocessed.bin\n", ref_dir.c_str()); - return 1; - } - fprintf(stderr, " Loaded preprocessed image: %d elements\n", ref_img.numel()); - - const int img_size = 1008; - bool ok = sam3_encode_image_from_preprocessed(*state, *model, ref_img.data.data(), img_size); - if (!ok) { - fprintf(stderr, "sam3_encode_image_from_preprocessed failed!\n"); - return 1; - } - fprintf(stderr, " Image encoded successfully\n"); - - // ═══ Collect all metrics ═══ - std::vector report; - - // Helper: dump state tensor, compare with ref in NHWC layout - auto check_nhwc = [&](const char * stage, const char * tensor_name, - const char * ref_name, float atol) { - sam3_dump_state_tensor(*state, tensor_name, cpp_out + "/" + tensor_name); - auto cpp = load_ref_f32(cpp_out + "/" + tensor_name); - auto ref = load_ref_f32(ref_dir + "/" + ref_name); - if (cpp.data.empty() || ref.data.empty()) { - fprintf(stderr, " [SKIP] %s\n", ref_name); - return; - } - int n = std::min(cpp.numel(), ref.numel()); - auto m = compute_full_metrics(cpp.data.data(), ref.data.data(), n, atol); - m.stage = stage; - m.name = ref_name; - // Build shape string from ref - std::string s; - for (size_t i = 0; i < ref.shape.size(); ++i) { - if (i > 0) s += ","; - s += std::to_string(ref.shape[i]); - } - m.py_shape_str = s; - report.push_back(m); - }; - - auto check_nchw = [&](const char * stage, const char * tensor_name, - const char * ref_name, float atol) { - sam3_dump_state_tensor(*state, tensor_name, cpp_out + "/" + tensor_name); - auto cpp = load_ref_f32(cpp_out + "/" + tensor_name); - auto ref = load_ref_f32(ref_dir + "/" + ref_name); - if (cpp.data.empty() || ref.data.empty()) { - fprintf(stderr, " [SKIP] %s\n", ref_name); - return; - } - // Transpose ggml→NCHW - if (cpp.shape.size() >= 3) { - int C = cpp.shape[0], W = cpp.shape[1], H = cpp.shape[2]; - std::vector transposed(C * W * H); - for (int c = 0; c < C; ++c) - for (int h = 0; h < H; ++h) - for (int w = 0; w < W; ++w) - transposed[c * H * W + h * W + w] = cpp.data[c + w * C + h * C * W]; - int n = std::min((int) transposed.size(), ref.numel()); - auto m = compute_full_metrics(transposed.data(), ref.data.data(), n, atol); - m.stage = stage; - m.name = ref_name; - std::string s; - for (size_t i = 0; i < ref.shape.size(); ++i) { - if (i > 0) s += ","; - s += std::to_string(ref.shape[i]); - } - m.py_shape_str = s; - report.push_back(m); - } - }; - - // ═══ Stage 2: ViT backbone ═══ - fprintf(stderr, "\n═══ Stage 2: ViT Backbone ═══\n"); - check_nhwc("ViT", "dbg_patch_embed", "patch_embed", 1e-4f); - check_nhwc("ViT", "dbg_after_pos_embed", "after_pos_embed", 1e-4f); - check_nhwc("ViT", "dbg_after_ln_pre", "after_ln_pre", 1e-4f); - - for (int i = 0; i < 32; ++i) { - char cpp_name[64], ref_name[64]; - snprintf(cpp_name, sizeof(cpp_name), "dbg_block_%d_out", i); - snprintf(ref_name, sizeof(ref_name), "block_%d_out", i); - auto ref = load_ref_f32(ref_dir + "/" + ref_name); - if (ref.data.empty()) continue; - // f32 accumulation across 32 blocks: max abs error ~0.06 on values up to - // magnitude ~280 (relative error < 0.02%). Cosine similarity is 1.0. - float atol = 0.07f; - check_nhwc("ViT", cpp_name, ref_name, atol); - } - - // ═══ Stage 3: Neck ═══ - fprintf(stderr, "\n═══ Stage 3: Neck (SimpleFPN) ═══\n"); - check_nchw("Neck", "vit_output", "vit_output_bchw", 0.05f); - for (int i = 0; i < 4; ++i) { - char name[32]; - snprintf(name, sizeof(name), "neck_det_%d", i); - check_nchw("Neck-Det", name, name, 5e-3f); - } - for (int i = 0; i < 4; ++i) { - char name[32]; - snprintf(name, sizeof(name), "neck_trk_%d", i); - auto ref = load_ref_f32(ref_dir + "/" + name); - if (ref.data.empty()) continue; - check_nchw("Neck-Trk", name, name, 5e-3f); - } - - // ═══ Stage 4-5: SAM Prompt Encoder + Mask Decoder ═══ - // Use sam3_test_dump_phase6_from_ref_inputs which loads neck data from reference - // files instead of reusing state tensors (avoids graph allocator buffer issues). - // This still tests the SAM decoder numerics end-to-end since the neck outputs - // were verified to match in Stage 3 above. - fprintf(stderr, "\n═══ Stage 4-5: SAM Prompt Encoder + Mask Decoder ═══\n"); - sam3_pvs_params pvs_params; - pvs_params.box = {200.0f, 50.0f, 800.0f, 950.0f}; - pvs_params.use_box = true; - pvs_params.multimask = false; - - if (!sam3_test_dump_phase6_from_ref_inputs(*model, ref_dir, pvs_params, - cpp_p6_out, params.n_threads)) { - fprintf(stderr, " [FAIL] sam3_test_dump_phase6_from_ref_inputs failed\n"); - } else { - const std::string p6_ref = ref_dir + "/phase6/cat_box"; - struct { const char * name; const char * label; float tol; } p6_tensors[] = { - {"sam_pe_sparse", "PE sparse embeddings", 1e-4f}, - {"sam_pe_dense", "PE dense embeddings", 1e-6f}, - {"sam_pe_image_pe", "PE image PE", 1e-5f}, - {"sam_dec_image_feats", "Dec image features", 5e-3f}, - {"sam_dec_tokens_initial", "Dec initial tokens", 1e-4f}, - {"sam_dec_block0_queries", "Dec block0 queries", 1e-4f}, - {"sam_dec_block0_keys", "Dec block0 keys", 1e-4f}, - {"sam_dec_block1_queries", "Dec block1 queries", 2e-4f}, - {"sam_dec_block1_keys", "Dec block1 keys", 2e-4f}, - {"sam_dec_final_queries", "Dec final queries", 2e-4f}, - {"sam_dec_feat_s1_proj", "Dec feat_s1 projection", 5e-3f}, - {"sam_dec_feat_s0_proj", "Dec feat_s0 projection", 5e-3f}, - {"sam_dec_upscaled", "Dec upscaled embedding", 9e-4f}, - {"sam_dec_mask_tokens", "Dec mask tokens", 2e-4f}, - {"sam_dec_masks", "Dec mask logits", 8e-3f}, - {"sam_dec_iou", "Dec IoU predictions", 2e-4f}, - {"sam_dec_obj_score", "Dec object score", 2e-4f}, - {"sam_dec_sam_token", "Dec SAM token", 2e-4f}, - }; - - for (const auto & t : p6_tensors) { - auto cpp = load_ref_f32(cpp_p6_out + "/" + t.name); - auto ref = load_ref_f32(p6_ref + "/" + t.name); - if (cpp.data.empty() || ref.data.empty()) { - fprintf(stderr, " [SKIP] %s (missing cpp=%s ref=%s)\n", - t.name, - cpp.data.empty() ? "yes" : "no", - ref.data.empty() ? "yes" : "no"); - continue; - } - int n = std::min(cpp.numel(), ref.numel()); - auto m = compute_full_metrics(cpp.data.data(), ref.data.data(), n, t.tol); - m.stage = "SAM-Dec"; - m.name = t.label; - std::string s; - for (size_t i = 0; i < ref.shape.size(); ++i) { - if (i > 0) s += ","; - s += std::to_string(ref.shape[i]); - } - m.py_shape_str = s; - report.push_back(m); - } - } - - // ═══ Print Full Report ═══ - fprintf(stderr, "\n"); - fprintf(stderr, "══════════════════════════════════════════════════════════════════════════════════════════════════════════════════════════\n"); - fprintf(stderr, " END-TO-END PVS PIPELINE COMPARISON: Python vs C++ (cat.jpg with box [200,50,800,950])\n"); - fprintf(stderr, "══════════════════════════════════════════════════════════════════════════════════════════════════════════════════════════\n\n"); - - int n_pass = 0, n_fail = 0; - fprintf(stderr, " %-8s %-6s %-30s %-16s %12s %12s %12s %12s %12s %12s\n", - "Status", "Stage", "Tensor", "Shape", "MAE", "Max", "RelErr", "Cosine", "P95", "P99"); - fprintf(stderr, " %-8s %-6s %-30s %-16s %12s %12s %12s %12s %12s %12s\n", - "------", "-----", "------", "-----", "---", "---", "------", "------", "---", "---"); - - for (const auto & m : report) { - const char * status = m.ok ? "PASS" : "FAIL"; - fprintf(stderr, " %-8s %-6s %-30s %-16s %12.4e %12.4e %12.4e %12.8f %12.4e %12.4e\n", - status, m.stage.c_str(), m.name.c_str(), m.py_shape_str.c_str(), - m.mae, m.max_abs, m.mean_rel, m.cosine, m.p95, m.p99); - if (m.ok) n_pass++; - else n_fail++; - } - - fprintf(stderr, "\n══════════════════════════════════════════════════════════════════════════════════════════════════════════════════════════\n"); - fprintf(stderr, " TOTAL: %d PASS, %d FAIL out of %d tensors\n", n_pass, n_fail, (int) report.size()); - fprintf(stderr, "══════════════════════════════════════════════════════════════════════════════════════════════════════════════════════════\n\n"); - - state.reset(); - sam3_free_model(*model); - model.reset(); - - return n_fail > 0 ? 1 : 0; -} diff --git a/tests/test_full_pipeline.py b/tests/test_full_pipeline.py deleted file mode 100644 index e9af497..0000000 --- a/tests/test_full_pipeline.py +++ /dev/null @@ -1,57 +0,0 @@ -#!/usr/bin/env python3 -"""Compare the full pipeline: C++ encode_image (with C++ preprocessing) -vs Python preprocessing → Python ViT. - -If the C++ bilinear resize matches PIL closely enough, the ViT outputs -should also be close. -""" -import numpy as np -import sys, os - -def load_tensor(path): - with open(path + ".shape") as f: - shape = [int(x) for x in f.read().strip().split(",") if x] - data = np.fromfile(path + ".bin", dtype=np.float32) - return data, shape - -def compare(name, ref, cpp, ref_shape, cpp_shape): - """Compare flat arrays, handling NCHW vs ggml [C, W, H] transposition.""" - # If shapes suggest NCHW vs ggml, transpose - if len(cpp_shape) == 3 and len(ref_shape) == 4: - C, W, H = cpp_shape - # ggml [C, W, H] flat: c + w*C + h*C*W - # NCHW flat: c*H*W + h*W + w - cpp_transposed = np.zeros(C * W * H, dtype=np.float32) - cpp_3d = cpp.reshape(H, W, C) # ggml column-major - # Transpose to [C, H, W] for NCHW comparison - cpp_nchw = cpp_3d.transpose(2, 0, 1).flatten() - cpp = cpp_nchw - ref = ref[:len(cpp)] - - diff = np.abs(ref - cpp) - max_d = diff.max() - mean_d = diff.mean() - cos = np.dot(ref, cpp) / (np.linalg.norm(ref) * np.linalg.norm(cpp) + 1e-12) - print(f" {name:40s} max={max_d:.6e} mean={mean_d:.6e} cos={cos:.8f}") - - -if __name__ == "__main__": - ref_dir = sys.argv[1] if len(sys.argv) > 1 else "tests/ref_phase3" - cpp_dir = sys.argv[2] if len(sys.argv) > 2 else "tests/cpp_out_full" - - if not os.path.exists(cpp_dir): - print(f"C++ output dir '{cpp_dir}' not found") - print("Run: ./build/tests/test_debug_encoder --full") - sys.exit(1) - - print("═══ Full Pipeline Comparison (C++ preprocess + ViT vs Python preprocess + ViT) ═══\n") - - for name in ["vit_output", "neck_det_0", "neck_det_1", "neck_det_2", "neck_det_3"]: - ref_path = ref_dir + "/" + ("vit_output_bchw" if name == "vit_output" else name) - cpp_path = cpp_dir + "/" + name - if not os.path.exists(ref_path + ".bin") or not os.path.exists(cpp_path + ".bin"): - print(f" [SKIP] {name}") - continue - ref, ref_shape = load_tensor(ref_path) - cpp, cpp_shape = load_tensor(cpp_path) - compare(name, ref, cpp, ref_shape, cpp_shape) diff --git a/tests/test_geom_enc.cpp b/tests/test_geom_enc.cpp deleted file mode 100644 index 478d78e..0000000 --- a/tests/test_geom_enc.cpp +++ /dev/null @@ -1,142 +0,0 @@ -#include "sam3.h" -#include "test_utils.h" - -#include -#include -#include - -struct geom_case { - std::string id; - sam3_pcs_params params; -}; - -int main(int argc, char** argv) { - if (argc < 4) { - fprintf(stderr, "Usage: %s \n", argv[0]); - return 1; - } - - const std::string ref_dir = argv[1]; - const std::string prephase_ref_dir = argv[2]; - const std::string model_path = argv[3]; - const std::string cpp_root = ref_dir + "/cpp_out_geom"; - - // Test cases matching the Python dump script - std::vector cases; - - // Case 1: dummy prompt (no exemplars, just CLS) - { - geom_case c; - c.id = "dummy_prompt"; - // No exemplars - cases.push_back(std::move(c)); - } - - // Case 2: single positive box (cx=0.5, cy=0.5, w=0.3, h=0.3) - // API uses XYXY format — convert from CxCyWH to XYXY - { - geom_case c; - c.id = "single_box"; - float cx = 0.5f, cy = 0.5f, w = 0.3f, h = 0.3f; - sam3_box box; - box.x0 = cx - w * 0.5f; // 0.35 - box.y0 = cy - h * 0.5f; // 0.35 - box.x1 = cx + w * 0.5f; // 0.65 - box.y1 = cy + h * 0.5f; // 0.65 - c.params.pos_exemplars.push_back(box); - cases.push_back(std::move(c)); - } - - // Case 3: two boxes (one positive, one negative) - // Box1: cx=0.3, cy=0.4, w=0.2, h=0.25 (positive) - // Box2: cx=0.7, cy=0.6, w=0.15, h=0.2 (negative) - { - geom_case c; - c.id = "two_boxes"; - { - float cx = 0.3f, cy = 0.4f, w = 0.2f, h = 0.25f; - sam3_box box = {cx - w * 0.5f, cy - h * 0.5f, cx + w * 0.5f, cy + h * 0.5f}; - c.params.pos_exemplars.push_back(box); - } - { - float cx = 0.7f, cy = 0.6f, w = 0.15f, h = 0.2f; - sam3_box box = {cx - w * 0.5f, cy - h * 0.5f, cx + w * 0.5f, cy + h * 0.5f}; - c.params.neg_exemplars.push_back(box); - } - cases.push_back(std::move(c)); - } - - sam3_params mparams; - mparams.model_path = model_path; - mparams.n_threads = 1; - mparams.use_gpu = false; - - auto model = sam3_load_model(mparams); - if (!model) { - fprintf(stderr, "Failed to load model\n"); - return 1; - } - - bool overall_ok = true; - float atol = 1e-4f; - - for (const auto& tc : cases) { - fprintf(stderr, "\n=== Geometry Encoder: %s ===\n", tc.id.c_str()); - - const std::string case_ref_dir = ref_dir + "/" + tc.id; - const std::string case_cpp_dir = cpp_root + "/" + tc.id; - ensure_dir(cpp_root); - ensure_dir(case_cpp_dir); - - if (!sam3_test_dump_geom_enc(*model, prephase_ref_dir, tc.params, - case_cpp_dir, mparams.n_threads)) { - fprintf(stderr, " [FAIL] dump failed for %s\n", tc.id.c_str()); - overall_ok = false; - continue; - } - - // Compare tensors - struct tensor_check { - std::string name; - float tol; - }; - std::vector checks = { - // post_final_proj is an input tensor whose buffer gets reused by the graph - // allocator during computation — comparing it after graph compute gives stale - // data. Only geom_output (the final transformer output) is meaningful. - {"geom_output", 5e-4f}, - }; - - for (const auto& chk : checks) { - auto ref = load_ref_f32(case_ref_dir + "/" + chk.name); - auto got = load_ref_f32(case_cpp_dir + "/" + chk.name); - - if (ref.data.empty() || got.data.empty()) { - fprintf(stderr, " [FAIL] %s: missing tensor\n", chk.name.c_str()); - overall_ok = false; - continue; - } - if (ref.numel() != got.numel()) { - fprintf(stderr, " [FAIL] %s: numel mismatch ref=%d got=%d\n", - chk.name.c_str(), ref.numel(), got.numel()); - overall_ok = false; - continue; - } - - auto r = compare_tensors(got.data.data(), ref.data.data(), ref.numel(), chk.tol); - bool ok = r.max_diff <= chk.tol; - fprintf(stderr, " [%s] %-30s max=%.3e mean=%.3e cos=%.8f tol=%.1e", - ok ? "PASS" : "FAIL", chk.name.c_str(), - r.max_diff, r.mean_diff, r.cosine_sim, chk.tol); - if (!ok) { - fprintf(stderr, " worst=%d got=%.6g ref=%.6g", - r.worst_index, r.worst_a, r.worst_b); - overall_ok = false; - } - fprintf(stderr, "\n"); - } - } - - fprintf(stderr, "\n%s\n", overall_ok ? "ALL PASS" : "SOME FAILED"); - return overall_ok ? 0 : 1; -} diff --git a/tests/test_jpeg_decode.cpp b/tests/test_jpeg_decode.cpp deleted file mode 100644 index b60d221..0000000 --- a/tests/test_jpeg_decode.cpp +++ /dev/null @@ -1,45 +0,0 @@ -// Compare stb_image JPEG decoding against PIL/libjpeg. -// Usage: ./test_jpeg_decode -#include "sam3.h" -#include "test_utils.h" -#include -#include -#include - -int main(int argc, char** argv) { - if (argc < 3) { - fprintf(stderr, "Usage: %s \n", argv[0]); - return 1; - } - auto img = sam3_load_image(argv[1]); - if (img.data.empty()) { fprintf(stderr, "Failed to load image\n"); return 1; } - fprintf(stderr, "stb_image decoded: %dx%d, %d bytes\n", - img.width, img.height, (int)img.data.size()); - - // Load PIL reference - std::ifstream f(argv[2], std::ios::binary); - if (!f) { fprintf(stderr, "Failed to load PIL ref\n"); return 1; } - std::vector pil_data(img.data.size()); - f.read(reinterpret_cast(pil_data.data()), pil_data.size()); - fprintf(stderr, "PIL decoded: %d bytes\n", (int)pil_data.size()); - - // Compare - int n_diff = 0, max_diff = 0; - for (size_t i = 0; i < img.data.size(); ++i) { - int d = abs((int)img.data[i] - (int)pil_data[i]); - if (d > 0) n_diff++; - if (d > max_diff) max_diff = d; - } - fprintf(stderr, "\n═══ JPEG Decoder Comparison ═══\n"); - fprintf(stderr, " max pixel diff: %d\n", max_diff); - fprintf(stderr, " pixels differing: %d / %d (%.2f%%)\n", - n_diff, (int)img.data.size(), 100.0 * n_diff / img.data.size()); - - // Check specific pixels - int y=16, x=517, c=2; - int idx = (y * img.width + x) * 3 + c; - fprintf(stderr, "\n stb pixel at (%d,%d,%d): %d\n", y, x, c, img.data[idx]); - fprintf(stderr, " PIL pixel at (%d,%d,%d): %d\n", y, x, c, pil_data[idx]); - - return max_diff > 0 ? 1 : 0; -} diff --git a/tests/test_llama_pcs.cpp b/tests/test_llama_pcs.cpp deleted file mode 100644 index 1f1cb90..0000000 --- a/tests/test_llama_pcs.cpp +++ /dev/null @@ -1,76 +0,0 @@ -/** - * Quick test: run PCS on the llama image with various text prompts. - * Uses sam3_encode_image (same as main_image example) — NOT preprocessed. - */ -#include "sam3.h" -#include -#include - -int main(int argc, char ** argv) { - if (argc < 3) { - fprintf(stderr, "Usage: %s [prompt]\n", argv[0]); - return 1; - } - - const std::string model_path = argv[1]; - const std::string image_path = argv[2]; - const std::string prompt = argc > 3 ? argv[3] : "llama"; - - sam3_params params; - params.model_path = model_path; - params.use_gpu = false; - params.n_threads = 4; - - fprintf(stderr, "Loading model...\n"); - auto model = sam3_load_model(params); - if (!model) { fprintf(stderr, "Failed to load model\n"); return 1; } - - auto state = sam3_create_state(*model, params); - if (!state) { fprintf(stderr, "Failed to create state\n"); return 1; } - - fprintf(stderr, "Loading image: %s\n", image_path.c_str()); - auto image = sam3_load_image(image_path); - if (image.data.empty()) { fprintf(stderr, "Failed to load image\n"); return 1; } - fprintf(stderr, "Image: %dx%d\n", image.width, image.height); - - fprintf(stderr, "Encoding image...\n"); - if (!sam3_encode_image(*state, *model, image)) { - fprintf(stderr, "Failed to encode image\n"); return 1; - } - - // Test multiple prompts - const char* prompts[] = {"llama", "guitar", "animal"}; - int n_prompts = 3; - if (argc > 4) { - prompts[0] = prompt.c_str(); - n_prompts = 1; - } - - for (int p = 0; p < n_prompts; ++p) { - fprintf(stderr, "\n═══ Testing prompt: '%s' ═══\n", prompts[p]); - - sam3_pcs_params pcs; - pcs.text_prompt = prompts[p]; - pcs.score_threshold = 0.3f; - pcs.nms_threshold = 0.1f; - - auto result = sam3_segment_pcs(*state, *model, pcs); - fprintf(stderr, "Result: %zu detections\n", result.detections.size()); - - for (size_t i = 0; i < result.detections.size() && i < 5; ++i) { - const auto& d = result.detections[i]; - fprintf(stderr, " det %zu: score=%.4f box=[%.1f,%.1f,%.1f,%.1f]\n", - i, d.score, d.box.x0, d.box.y0, d.box.x1, d.box.y1); - } - - // Save first mask - if (!result.detections.empty()) { - std::string mask_path = "/tmp/llama_mask_" + std::string(prompts[p]) + ".png"; - sam3_save_mask(result.detections[0].mask, mask_path); - fprintf(stderr, " Saved mask: %s\n", mask_path.c_str()); - } - } - - sam3_free_model(*model); - return 0; -} diff --git a/tests/test_load.cpp b/tests/test_load.cpp deleted file mode 100644 index f9d0b04..0000000 --- a/tests/test_load.cpp +++ /dev/null @@ -1,38 +0,0 @@ -#include "sam3.h" -#include - -int main(int argc, char ** argv) { - if (argc < 2) { - fprintf(stderr, "Usage: %s \n", argv[0]); - return 1; - } - - sam3_params params; - params.model_path = argv[1]; - params.use_gpu = true; - params.n_threads = 4; - - auto model = sam3_load_model(params); - if (!model) { - fprintf(stderr, "Failed to load model\n"); - return 1; - } - - fprintf(stderr, "Model loaded successfully!\n"); - - auto state = sam3_create_state(*model, params); - if (!state) { - fprintf(stderr, "Failed to create state\n"); - return 1; - } - - fprintf(stderr, "State created successfully!\n"); - - // Explicit cleanup before ggml global shutdown - state.reset(); - sam3_free_model(*model); - model.reset(); - - fprintf(stderr, "Cleanup complete.\n"); - return 0; -} diff --git a/tests/test_mem_enc.cpp b/tests/test_mem_enc.cpp deleted file mode 100644 index 2d35ea3..0000000 --- a/tests/test_mem_enc.cpp +++ /dev/null @@ -1,110 +0,0 @@ -/** - * Memory encoder comparison test. - * Encodes the llama image, runs PVS with a point, then dumps the - * memory encoder outputs for comparison against Python reference. - * - * Usage: test_mem_enc [x y] - */ -#include "sam3.h" -#include "test_utils.h" -#include -#include - -int main(int argc, char ** argv) { - if (argc < 3) { - fprintf(stderr, "Usage: %s [x y]\n", argv[0]); - return 1; - } - - const std::string model_path = argv[1]; - const std::string ref_dir = std::string(argv[2]) + "/video_ref"; - float px = argc > 4 ? atof(argv[3]) : 315.0f; - float py = argc > 4 ? atof(argv[4]) : 250.0f; - - sam3_params params; - params.model_path = model_path; - params.use_gpu = false; - params.n_threads = 4; - - auto model = sam3_load_model(params); - if (!model) { fprintf(stderr, "Failed to load model\n"); return 1; } - - auto state = sam3_create_state(*model, params); - if (!state) { fprintf(stderr, "Failed to create state\n"); return 1; } - - // Load llama image - auto image = sam3_load_image("../ggml/examples/sam/example.jpg"); - if (image.data.empty()) { fprintf(stderr, "Failed to load image\n"); return 1; } - - // Encode image - if (!sam3_encode_image(*state, *model, image)) { - fprintf(stderr, "Failed to encode image\n"); return 1; - } - - // Create tracker - sam3_video_params vp; - vp.hotstart_delay = 0; - auto tracker = sam3_create_tracker(*model, vp); - - // Add instance with point prompt → triggers PVS + memory encoding - sam3_pvs_params pvs; - pvs.pos_points.push_back({px, py}); - pvs.multimask = false; - - int inst_id = sam3_tracker_add_instance(*tracker, *state, *model, pvs); - fprintf(stderr, "Added instance %d\n", inst_id); - - if (inst_id < 0) { - fprintf(stderr, "Failed to add instance\n"); - return 1; - } - - // ── Compare memory encoder outputs against Python reference ────── - fprintf(stderr, "\n═══ Memory Encoder Comparison ═══\n"); - - // Note: We're comparing C++ (stb_image decoded) vs Python (PIL decoded). - // The images differ slightly, so we expect some divergence. - // We check that the outputs are in the same range and structure. - - // Check that Python reference exists - auto py_mem_feat = load_ref_f32(ref_dir + "/mem_enc_output_features"); - auto py_mem_pe = load_ref_f32(ref_dir + "/mem_enc_output_pe_0"); - auto py_obj_ptr = load_ref_f32(ref_dir + "/obj_ptr_output"); - - fprintf(stderr, "Python reference:\n"); - if (!py_mem_feat.data.empty()) - fprintf(stderr, " mem_enc_output: numel=%d\n", py_mem_feat.numel()); - if (!py_mem_pe.data.empty()) - fprintf(stderr, " mem_enc_pe: numel=%d\n", py_mem_pe.numel()); - if (!py_obj_ptr.data.empty()) - fprintf(stderr, " obj_ptr: numel=%d\n", py_obj_ptr.numel()); - - // The C++ memory encoder stored its output in the tracker's memory bank. - // We can't easily extract it from the tracker. But we can verify the - // pipeline ran correctly by checking the tracker state. - fprintf(stderr, "\nC++ tracker state:\n"); - fprintf(stderr, " frame_index: %d\n", sam3_tracker_frame_index(*tracker)); - - // Dump Python reference ranges for comparison - if (!py_mem_feat.data.empty()) { - float mn = py_mem_feat.data[0], mx = py_mem_feat.data[0]; - for (float v : py_mem_feat.data) { mn = std::min(mn, v); mx = std::max(mx, v); } - fprintf(stderr, " Python mem_enc_features range: [%.4f, %.4f]\n", mn, mx); - } - - if (!py_obj_ptr.data.empty()) { - float mn = py_obj_ptr.data[0], mx = py_obj_ptr.data[0]; - for (float v : py_obj_ptr.data) { mn = std::min(mn, v); mx = std::max(mx, v); } - fprintf(stderr, " Python obj_ptr range: [%.4f, %.4f]\n", mn, mx); - fprintf(stderr, " Python obj_ptr[0..4]: [%.4f, %.4f, %.4f, %.4f, %.4f]\n", - py_obj_ptr.data[0], py_obj_ptr.data[1], py_obj_ptr.data[2], - py_obj_ptr.data[3], py_obj_ptr.data[4]); - } - - fprintf(stderr, "\n═══ Result: Memory encoder reference data available ═══\n"); - fprintf(stderr, "For proper numerical comparison, use sam3_encode_image_from_preprocessed()\n"); - fprintf(stderr, "with the Python-exported preprocessed image to eliminate JPEG decoder differences.\n"); - - sam3_free_model(*model); - return 0; -} diff --git a/tests/test_metal_block0_stage_compare.cpp b/tests/test_metal_block0_stage_compare.cpp deleted file mode 100644 index a8a9c9a..0000000 --- a/tests/test_metal_block0_stage_compare.cpp +++ /dev/null @@ -1,398 +0,0 @@ -#include "sam3.h" -#include "test_utils.h" - -#include -#include -#include -#include -#include - -struct dump_item { - const char * label; - const char * name; -}; - -static std::string format_shape(const std::vector & shape) { - std::string out = "["; - for (size_t i = 0; i < shape.size(); ++i) { - if (i > 0) { - out += ", "; - } - out += std::to_string(shape[i]); - } - out += "]"; - return out; -} - -static std::vector trim_shape(const int64_t ne[4]) { - std::vector shape = { - (int) ne[0], - (int) ne[1], - (int) ne[2], - (int) ne[3], - }; - while (!shape.empty() && shape.back() == 1) { - shape.pop_back(); - } - if (shape.empty()) { - shape.push_back(1); - } - return shape; -} - -static bool dump_named_tensors(const sam3_state & state, - const std::vector & items, - const std::string & output_dir) { - for (const auto & item : items) { - if (!sam3_dump_state_tensor(state, item.name, output_dir + "/" + item.name)) { - fprintf(stderr, "failed to dump %s\n", item.name); - return false; - } - } - return true; -} - -static bool run_selective_vit(sam3_state & state, - const sam3_model & model, - const std::vector & chw_data, - int img_size, - const std::vector & items, - const std::string & output_dir) { - std::vector names; - names.reserve(items.size()); - for (const auto & item : items) { - names.push_back(item.name); - } - - if (!sam3_encode_vit_from_preprocessed_selective(state, model, chw_data.data(), img_size, names)) { - return false; - } - return dump_named_tensors(state, items, output_dir); -} - -static ref_tensor_f32 run_stage_or_die(const sam3_model & model, - sam3_vit_block_stage stage, - const ref_tensor_f32 & input, - int n_threads) { - int64_t in_ne[4] = {1, 1, 1, 1}; - for (size_t i = 0; i < input.shape.size() && i < 4; ++i) { - in_ne[i] = input.shape[i]; - } - - int64_t out_ne[4] = {0, 0, 0, 0}; - ref_tensor_f32 out; - if (!sam3_test_run_vit_block_stage(model, 0, stage, input.data.data(), in_ne, out.data, out_ne, n_threads)) { - fprintf(stderr, "stage %d failed\n", (int) stage); - std::exit(1); - } - out.shape = trim_shape(out_ne); - return out; -} - -static ref_tensor_f32 run_linear_host_ref_or_die(const sam3_model & model, - sam3_vit_block_stage stage, - const ref_tensor_f32 & input, - bool use_double_accum) { - int64_t in_ne[4] = {1, 1, 1, 1}; - for (size_t i = 0; i < input.shape.size() && i < 4; ++i) { - in_ne[i] = input.shape[i]; - } - - int64_t out_ne[4] = {0, 0, 0, 0}; - ref_tensor_f32 out; - if (!sam3_test_run_vit_block_linear_host_ref(model, 0, stage, input.data.data(), in_ne, use_double_accum, out.data, out_ne)) { - fprintf(stderr, "host linear ref stage %d failed\n", (int) stage); - std::exit(1); - } - out.shape = trim_shape(out_ne); - return out; -} - -static ref_tensor_f32 run_block0_input_or_die(const sam3_model & model, - const std::vector & chw_data, - int img_size, - int n_threads) { - int64_t out_ne[4] = {0, 0, 0, 0}; - ref_tensor_f32 out; - if (!sam3_test_run_vit_block0_input(model, chw_data.data(), img_size, out.data, out_ne, n_threads)) { - fprintf(stderr, "block0 input prefix failed\n"); - std::exit(1); - } - out.shape = trim_shape(out_ne); - return out; -} - -static ref_tensor_f32 add_ref(const ref_tensor_f32 & a, const ref_tensor_f32 & b) { - if (a.shape != b.shape || a.data.size() != b.data.size()) { - fprintf(stderr, "add_ref: shape mismatch\n"); - std::exit(1); - } - ref_tensor_f32 out; - out.shape = a.shape; - out.data.resize(a.data.size()); - for (size_t i = 0; i < a.data.size(); ++i) { - out.data[i] = a.data[i] + b.data[i]; - } - return out; -} - -static void print_diff_row(const char * label, - const ref_tensor_f32 & a, - const ref_tensor_f32 & b, - float atol = 1e-4f) { - const compare_result r = compare_tensors(a.data.data(), b.data.data(), a.numel(), atol); - fprintf(stderr, "%-22s %-18s %12.6f %12.6f %10d\n", - label, - format_shape(a.shape).c_str(), - r.max_diff, - r.mean_diff, - r.n_bad); -} - -static void print_graph_row(const dump_item & item, - const std::string & cpu_dir, - const std::string & metal_dir) { - const ref_tensor_f32 cpu = load_ref_f32(cpu_dir + "/" + item.name); - const ref_tensor_f32 metal = load_ref_f32(metal_dir + "/" + item.name); - if (cpu.data.empty() || metal.data.empty()) { - fprintf(stderr, "%-22s %-18s %12s %12s %10s\n", - item.label, "-", "load-fail", "load-fail", "-"); - return; - } - print_diff_row(item.label, cpu, metal); -} - -int main(int argc, char ** argv) { - if (argc < 2) { - fprintf(stderr, "Usage: %s [ref_dir]\n", argv[0]); - fprintf(stderr, "Default ref_dir: tests/ref_phase3\n"); - return 1; - } - - const std::string model_path = argv[1]; - const std::string ref_dir = argc >= 3 ? argv[2] : "tests/ref_phase3"; - - const std::vector graph_items = { - { "block_input", "dbg_after_ln_pre" }, - { "qkv_proj_output", "dbg_block_0_qkv_proj"}, - { "attn_proj_output", "dbg_block_0_attn_proj"}, - { "resid1_output", "dbg_block_0_resid1" }, - { "mlp_output", "dbg_block_0_mlp" }, - { "block_output", "dbg_block_0_out" }, - }; - - const std::string cpu_dir = "/tmp/sam3_block0_stage_cpu"; - const std::string metal_dir = "/tmp/sam3_block0_stage_metal"; - ensure_dir(cpu_dir); - ensure_dir(metal_dir); - - const ref_tensor_f32 preprocessed = load_ref_f32(ref_dir + "/preprocessed"); - if (preprocessed.data.empty() || preprocessed.shape.size() != 4) { - fprintf(stderr, "failed to load %s/preprocessed\n", ref_dir.c_str()); - return 1; - } - const int img_size = preprocessed.shape[2]; - const int n_threads = 8; - - sam3_params cpu_params; - cpu_params.model_path = model_path; - cpu_params.use_gpu = false; - cpu_params.n_threads = n_threads; - - sam3_params metal_params = cpu_params; - metal_params.use_gpu = true; - - auto cpu_model = sam3_load_model(cpu_params); - auto metal_model = sam3_load_model(metal_params); - if (!cpu_model || !metal_model) { - fprintf(stderr, "failed to load CPU or Metal model\n"); - return 1; - } - - auto cpu_state = sam3_create_state(*cpu_model, cpu_params); - auto metal_state = sam3_create_state(*metal_model, metal_params); - if (!cpu_state || !metal_state) { - fprintf(stderr, "failed to create CPU or Metal state\n"); - return 1; - } - - if (!run_selective_vit(*cpu_state, *cpu_model, preprocessed.data, img_size, graph_items, cpu_dir)) { - fprintf(stderr, "CPU selective block-0 run failed\n"); - return 1; - } - if (!run_selective_vit(*metal_state, *metal_model, preprocessed.data, img_size, graph_items, metal_dir)) { - fprintf(stderr, "Metal selective block-0 run failed\n"); - return 1; - } - - fprintf(stderr, "\n=== Block-0 graph outputs: CPU vs Metal ===\n"); - fprintf(stderr, "%-22s %-18s %12s %12s %10s\n", - "checkpoint", "shape", "max_abs_diff", "mean_abs", "n_bad"); - for (const auto & item : graph_items) { - print_graph_row(item, cpu_dir, metal_dir); - } - - const ref_tensor_f32 cpu_block_input = run_block0_input_or_die(*cpu_model, preprocessed.data, img_size, n_threads); - const ref_tensor_f32 metal_block_input = run_block0_input_or_die(*metal_model, preprocessed.data, img_size, n_threads); - - sam3_tensor_info qkv_info = {}; - sam3_tensor_info proj_info = {}; - sam3_tensor_info mlp1_info = {}; - sam3_tensor_info rope_info = {}; - sam3_get_model_tensor_info(*cpu_model, "vit.blocks.0.attn.qkv.weight", qkv_info); - sam3_get_model_tensor_info(*cpu_model, "vit.blocks.0.attn.proj.weight", proj_info); - sam3_get_model_tensor_info(*cpu_model, "vit.blocks.0.mlp.lin1.weight", mlp1_info); - sam3_get_model_tensor_info(*cpu_model, "vit.blocks.0.attn.freqs_cis", rope_info); - - fprintf(stderr, "\n=== Exact block-0 case ===\n"); - fprintf(stderr, "block input shape: %s\n", format_shape(cpu_block_input.shape).c_str()); - fprintf(stderr, "window attention: full_spatial=72x72 window=24 heads=%d head_dim=%d\n", 16, 64); - fprintf(stderr, "qkv weight: shape=[%lld,%lld] type=%d\n", - (long long) qkv_info.ne[0], (long long) qkv_info.ne[1], qkv_info.type); - fprintf(stderr, "proj weight: shape=[%lld,%lld] type=%d\n", - (long long) proj_info.ne[0], (long long) proj_info.ne[1], proj_info.type); - fprintf(stderr, "mlp fc1: shape=[%lld,%lld] type=%d\n", - (long long) mlp1_info.ne[0], (long long) mlp1_info.ne[1], mlp1_info.type); - fprintf(stderr, "freqs_cis: shape=[%lld,%lld,%lld] type=%d\n", - (long long) rope_info.ne[0], (long long) rope_info.ne[1], (long long) rope_info.ne[2], rope_info.type); - print_diff_row("block0_input_exact", cpu_block_input, metal_block_input); - - const ref_tensor_f32 cpu_norm1_self = run_stage_or_die(*cpu_model, SAM3_VIT_BLOCK_STAGE_NORM1, cpu_block_input, n_threads); - const ref_tensor_f32 metal_norm1_self = run_stage_or_die(*metal_model, SAM3_VIT_BLOCK_STAGE_NORM1, metal_block_input, n_threads); - const ref_tensor_f32 cpu_win_part_self = run_stage_or_die(*cpu_model, SAM3_VIT_BLOCK_STAGE_WINDOW_PART, cpu_norm1_self, n_threads); - const ref_tensor_f32 metal_win_part_self = run_stage_or_die(*metal_model, SAM3_VIT_BLOCK_STAGE_WINDOW_PART, metal_norm1_self, n_threads); - const ref_tensor_f32 cpu_qkv_self = run_stage_or_die(*cpu_model, SAM3_VIT_BLOCK_STAGE_QKV_PROJ, cpu_win_part_self, n_threads); - const ref_tensor_f32 metal_qkv_self = run_stage_or_die(*metal_model, SAM3_VIT_BLOCK_STAGE_QKV_PROJ, metal_win_part_self, n_threads); - const ref_tensor_f32 cpu_attn_self = run_stage_or_die(*cpu_model, SAM3_VIT_BLOCK_STAGE_ATTN_CORE, cpu_qkv_self, n_threads); - const ref_tensor_f32 metal_attn_self = run_stage_or_die(*metal_model, SAM3_VIT_BLOCK_STAGE_ATTN_CORE, metal_qkv_self, n_threads); - const ref_tensor_f32 cpu_attn_proj_self = run_stage_or_die(*cpu_model, SAM3_VIT_BLOCK_STAGE_ATTN_PROJ, cpu_attn_self, n_threads); - const ref_tensor_f32 metal_attn_proj_self = run_stage_or_die(*metal_model, SAM3_VIT_BLOCK_STAGE_ATTN_PROJ, metal_attn_self, n_threads); - const ref_tensor_f32 cpu_unpart_self = run_stage_or_die(*cpu_model, SAM3_VIT_BLOCK_STAGE_WINDOW_UNPART, cpu_attn_proj_self, n_threads); - const ref_tensor_f32 metal_unpart_self = run_stage_or_die(*metal_model, SAM3_VIT_BLOCK_STAGE_WINDOW_UNPART, metal_attn_proj_self, n_threads); - const ref_tensor_f32 cpu_resid1_self = add_ref(cpu_block_input, cpu_unpart_self); - const ref_tensor_f32 metal_resid1_self = add_ref(metal_block_input, metal_unpart_self); - const ref_tensor_f32 cpu_norm2_self = run_stage_or_die(*cpu_model, SAM3_VIT_BLOCK_STAGE_NORM2, cpu_resid1_self, n_threads); - const ref_tensor_f32 metal_norm2_self = run_stage_or_die(*metal_model, SAM3_VIT_BLOCK_STAGE_NORM2, metal_resid1_self, n_threads); - const ref_tensor_f32 cpu_mlp_self = run_stage_or_die(*cpu_model, SAM3_VIT_BLOCK_STAGE_MLP, cpu_norm2_self, n_threads); - const ref_tensor_f32 metal_mlp_self = run_stage_or_die(*metal_model, SAM3_VIT_BLOCK_STAGE_MLP, metal_norm2_self, n_threads); - const ref_tensor_f32 cpu_block_self = add_ref(cpu_resid1_self, cpu_mlp_self); - const ref_tensor_f32 metal_block_self = add_ref(metal_resid1_self, metal_mlp_self); - - const ref_tensor_f32 cpu_norm1_shared = run_stage_or_die(*cpu_model, SAM3_VIT_BLOCK_STAGE_NORM1, cpu_block_input, n_threads); - const ref_tensor_f32 metal_norm1_shared = run_stage_or_die(*metal_model, SAM3_VIT_BLOCK_STAGE_NORM1, cpu_block_input, n_threads); - const ref_tensor_f32 cpu_win_part_shared = run_stage_or_die(*cpu_model, SAM3_VIT_BLOCK_STAGE_WINDOW_PART, cpu_norm1_shared, n_threads); - const ref_tensor_f32 metal_win_part_shared = run_stage_or_die(*metal_model, SAM3_VIT_BLOCK_STAGE_WINDOW_PART, cpu_norm1_shared, n_threads); - const ref_tensor_f32 cpu_qkv_shared = run_stage_or_die(*cpu_model, SAM3_VIT_BLOCK_STAGE_QKV_PROJ, cpu_win_part_shared, n_threads); - const ref_tensor_f32 metal_qkv_shared = run_stage_or_die(*metal_model, SAM3_VIT_BLOCK_STAGE_QKV_PROJ, cpu_win_part_shared, n_threads); - const ref_tensor_f32 cpu_attn_shared = run_stage_or_die(*cpu_model, SAM3_VIT_BLOCK_STAGE_ATTN_CORE, cpu_qkv_shared, n_threads); - const ref_tensor_f32 metal_attn_shared = run_stage_or_die(*metal_model, SAM3_VIT_BLOCK_STAGE_ATTN_CORE, cpu_qkv_shared, n_threads); - const ref_tensor_f32 cpu_attn_proj_shared = run_stage_or_die(*cpu_model, SAM3_VIT_BLOCK_STAGE_ATTN_PROJ, cpu_attn_shared, n_threads); - const ref_tensor_f32 metal_attn_proj_shared = run_stage_or_die(*metal_model, SAM3_VIT_BLOCK_STAGE_ATTN_PROJ, cpu_attn_shared, n_threads); - const ref_tensor_f32 cpu_unpart_shared = run_stage_or_die(*cpu_model, SAM3_VIT_BLOCK_STAGE_WINDOW_UNPART, cpu_attn_proj_shared, n_threads); - const ref_tensor_f32 metal_unpart_shared = run_stage_or_die(*metal_model, SAM3_VIT_BLOCK_STAGE_WINDOW_UNPART, cpu_attn_proj_shared, n_threads); - const ref_tensor_f32 cpu_resid1_shared = add_ref(cpu_block_input, cpu_unpart_shared); - const ref_tensor_f32 metal_resid1_shared = add_ref(cpu_block_input, metal_unpart_shared); - const ref_tensor_f32 cpu_norm2_shared = run_stage_or_die(*cpu_model, SAM3_VIT_BLOCK_STAGE_NORM2, cpu_resid1_shared, n_threads); - const ref_tensor_f32 metal_norm2_shared = run_stage_or_die(*metal_model, SAM3_VIT_BLOCK_STAGE_NORM2, cpu_resid1_shared, n_threads); - const ref_tensor_f32 cpu_mlp_fc1_shared = run_stage_or_die(*cpu_model, SAM3_VIT_BLOCK_STAGE_MLP_FC1, cpu_norm2_shared, n_threads); - const ref_tensor_f32 metal_mlp_fc1_shared = run_stage_or_die(*metal_model, SAM3_VIT_BLOCK_STAGE_MLP_FC1, cpu_norm2_shared, n_threads); - const ref_tensor_f32 cpu_mlp_gelu_shared = run_stage_or_die(*cpu_model, SAM3_VIT_BLOCK_STAGE_MLP_GELU, cpu_mlp_fc1_shared, n_threads); - const ref_tensor_f32 metal_mlp_gelu_shared = run_stage_or_die(*metal_model, SAM3_VIT_BLOCK_STAGE_MLP_GELU, cpu_mlp_fc1_shared, n_threads); - const ref_tensor_f32 cpu_mlp_fc2_shared = run_stage_or_die(*cpu_model, SAM3_VIT_BLOCK_STAGE_MLP_FC2, cpu_mlp_gelu_shared, n_threads); - const ref_tensor_f32 metal_mlp_fc2_shared = run_stage_or_die(*metal_model, SAM3_VIT_BLOCK_STAGE_MLP_FC2, cpu_mlp_gelu_shared, n_threads); - const ref_tensor_f32 cpu_mlp_shared = run_stage_or_die(*cpu_model, SAM3_VIT_BLOCK_STAGE_MLP, cpu_norm2_shared, n_threads); - const ref_tensor_f32 metal_mlp_shared = run_stage_or_die(*metal_model, SAM3_VIT_BLOCK_STAGE_MLP, cpu_norm2_shared, n_threads); - const ref_tensor_f32 cpu_block_shared = add_ref(cpu_resid1_shared, cpu_mlp_shared); - const ref_tensor_f32 metal_block_shared = add_ref(cpu_resid1_shared, metal_mlp_shared); - - const std::map cpu_graph = { - {"qkv_proj_output", load_ref_f32(cpu_dir + "/dbg_block_0_qkv_proj")}, - {"attn_proj_output", load_ref_f32(cpu_dir + "/dbg_block_0_attn_proj")}, - {"resid1_output", load_ref_f32(cpu_dir + "/dbg_block_0_resid1")}, - {"mlp_output", load_ref_f32(cpu_dir + "/dbg_block_0_mlp")}, - {"block_output", load_ref_f32(cpu_dir + "/dbg_block_0_out")}, - }; - const std::map metal_graph = { - {"qkv_proj_output", load_ref_f32(metal_dir + "/dbg_block_0_qkv_proj")}, - {"attn_proj_output", load_ref_f32(metal_dir + "/dbg_block_0_attn_proj")}, - {"resid1_output", load_ref_f32(metal_dir + "/dbg_block_0_resid1")}, - {"mlp_output", load_ref_f32(metal_dir + "/dbg_block_0_mlp")}, - {"block_output", load_ref_f32(metal_dir + "/dbg_block_0_out")}, - }; - - fprintf(stderr, "\n=== Exact isolated stages on shared canonical CPU inputs ===\n"); - fprintf(stderr, "%-22s %-18s %12s %12s %10s\n", - "stage", "shape", "max_abs_diff", "mean_abs", "n_bad"); - print_diff_row("norm1_output", cpu_norm1_shared, metal_norm1_shared); - print_diff_row("win_part_output", cpu_win_part_shared, metal_win_part_shared); - print_diff_row("qkv_proj_output", cpu_qkv_shared, metal_qkv_shared); - print_diff_row("attn_core_output", cpu_attn_shared, metal_attn_shared); - print_diff_row("attn_proj_output", cpu_attn_proj_shared, metal_attn_proj_shared); - print_diff_row("win_unpart_output", cpu_unpart_shared, metal_unpart_shared); - print_diff_row("resid1_output", cpu_resid1_shared, metal_resid1_shared); - print_diff_row("norm2_output", cpu_norm2_shared, metal_norm2_shared); - print_diff_row("mlp_output", cpu_mlp_shared, metal_mlp_shared); - print_diff_row("block_output", cpu_block_shared, metal_block_shared); - - const ref_tensor_f32 host_qkv_f64 = run_linear_host_ref_or_die(*cpu_model, SAM3_VIT_BLOCK_STAGE_QKV_PROJ, cpu_win_part_shared, true); - const ref_tensor_f32 host_attn_proj_f64 = run_linear_host_ref_or_die(*cpu_model, SAM3_VIT_BLOCK_STAGE_ATTN_PROJ, cpu_attn_shared, true); - const ref_tensor_f32 host_mlp_fc1_f64 = run_linear_host_ref_or_die(*cpu_model, SAM3_VIT_BLOCK_STAGE_MLP_FC1, cpu_norm2_shared, true); - const ref_tensor_f32 host_mlp_fc2_f64 = run_linear_host_ref_or_die(*cpu_model, SAM3_VIT_BLOCK_STAGE_MLP_FC2, cpu_mlp_gelu_shared, true); - - fprintf(stderr, "\n=== Block-0 host linear refs (double accumulation) ===\n"); - fprintf(stderr, "%-22s %-18s %12s %12s %10s\n", - "stage", "shape", "max_abs_diff", "mean_abs", "n_bad"); - print_diff_row("qkv_cpu_vs_host", cpu_qkv_shared, host_qkv_f64); - print_diff_row("qkv_mtl_vs_host", metal_qkv_shared, host_qkv_f64); - print_diff_row("attn_proj_cpu_vs_host", cpu_attn_proj_shared, host_attn_proj_f64); - print_diff_row("attn_proj_mtl_vs_host", metal_attn_proj_shared, host_attn_proj_f64); - print_diff_row("mlp_fc1_cpu_vs_host", cpu_mlp_fc1_shared, host_mlp_fc1_f64); - print_diff_row("mlp_fc1_mtl_vs_host", metal_mlp_fc1_shared, host_mlp_fc1_f64); - print_diff_row("mlp_fc2_cpu_vs_host", cpu_mlp_fc2_shared, host_mlp_fc2_f64); - print_diff_row("mlp_fc2_mtl_vs_host", metal_mlp_fc2_shared, host_mlp_fc2_f64); - - fprintf(stderr, "\n=== Block-0 MLP split on shared canonical CPU inputs ===\n"); - fprintf(stderr, "%-22s %-18s %12s %12s %10s\n", - "stage", "shape", "max_abs_diff", "mean_abs", "n_bad"); - print_diff_row("mlp_fc1_output", cpu_mlp_fc1_shared, metal_mlp_fc1_shared); - print_diff_row("mlp_gelu_output", cpu_mlp_gelu_shared, metal_mlp_gelu_shared); - print_diff_row("mlp_fc2_output", cpu_mlp_fc2_shared, metal_mlp_fc2_shared); - - fprintf(stderr, "\n=== Isolated stage self-consistency vs graph outputs ===\n"); - fprintf(stderr, "%-22s %-18s %-24s %-24s\n", - "stage", "shape", "cpu_iso_vs_cpu_graph", "metal_iso_vs_metal_graph"); - - auto print_self_row = [&](const char * label, - const ref_tensor_f32 & cpu_iso, - const ref_tensor_f32 & metal_iso, - const char * graph_key) { - const auto cpu_it = cpu_graph.find(graph_key); - const auto metal_it = metal_graph.find(graph_key); - if (cpu_it == cpu_graph.end() || metal_it == metal_graph.end()) { - fprintf(stderr, "%-22s %-18s %-24s %-24s\n", - label, "-", "missing", "missing"); - return; - } - - const compare_result cpu_r = compare_tensors(cpu_iso.data.data(), cpu_it->second.data.data(), cpu_iso.numel(), 1e-4f); - const compare_result metal_r = compare_tensors(metal_iso.data.data(), metal_it->second.data.data(), metal_iso.numel(), 1e-4f); - - char cpu_buf[64]; - char metal_buf[64]; - snprintf(cpu_buf, sizeof(cpu_buf), "max=%.6f mean=%.6f", cpu_r.max_diff, cpu_r.mean_diff); - snprintf(metal_buf, sizeof(metal_buf), "max=%.6f mean=%.6f", metal_r.max_diff, metal_r.mean_diff); - fprintf(stderr, "%-22s %-18s %-24s %-24s\n", - label, - format_shape(cpu_iso.shape).c_str(), - cpu_buf, - metal_buf); - }; - - print_self_row("qkv_proj_output", cpu_qkv_self, metal_qkv_self, "qkv_proj_output"); - print_self_row("attn_proj_output", cpu_attn_proj_self, metal_attn_proj_self, "attn_proj_output"); - print_self_row("resid1_output", cpu_resid1_self, metal_resid1_self, "resid1_output"); - print_self_row("mlp_output", cpu_mlp_self, metal_mlp_self, "mlp_output"); - print_self_row("block_output", cpu_block_self, metal_block_self, "block_output"); - - return 0; -} diff --git a/tests/test_metal_block_stage_compare.cpp b/tests/test_metal_block_stage_compare.cpp deleted file mode 100644 index 28c486e..0000000 --- a/tests/test_metal_block_stage_compare.cpp +++ /dev/null @@ -1,535 +0,0 @@ -#include "sam3.h" -#include "test_utils.h" - -#include -#include -#include -#include -#include - -struct dump_item { - const char * label; - std::string name; -}; - -static std::string format_shape(const std::vector & shape) { - std::string out = "["; - for (size_t i = 0; i < shape.size(); ++i) { - if (i > 0) { - out += ", "; - } - out += std::to_string(shape[i]); - } - out += "]"; - return out; -} - -static std::vector trim_shape(const int64_t ne[4]) { - std::vector shape = { - (int) ne[0], - (int) ne[1], - (int) ne[2], - (int) ne[3], - }; - while (!shape.empty() && shape.back() == 1) { - shape.pop_back(); - } - if (shape.empty()) { - shape.push_back(1); - } - return shape; -} - -static std::string make_block_name(int block_idx, const char * suffix) { - char buf[64]; - snprintf(buf, sizeof(buf), "dbg_block_%d_%s", block_idx, suffix); - return buf; -} - -static bool dump_named_tensors(const sam3_state & state, - const std::vector & items, - const std::string & output_dir) { - for (const auto & item : items) { - if (!sam3_dump_state_tensor(state, item.name, output_dir + "/" + item.name)) { - fprintf(stderr, "failed to dump %s\n", item.name.c_str()); - return false; - } - } - return true; -} - -static bool run_selective_vit(sam3_state & state, - const sam3_model & model, - const std::vector & chw_data, - int img_size, - int debug_block_idx, - const std::vector & items, - const std::string & output_dir) { - std::vector names; - names.reserve(items.size()); - for (const auto & item : items) { - names.push_back(item.name); - } - - if (!sam3_encode_vit_from_preprocessed_selective(state, model, chw_data.data(), img_size, names)) { - return false; - } - return dump_named_tensors(state, items, output_dir); -} - -static ref_tensor_f32 run_stage_or_die(const sam3_model & model, - int block_idx, - sam3_vit_block_stage stage, - const ref_tensor_f32 & input, - int n_threads) { - int64_t in_ne[4] = {1, 1, 1, 1}; - for (size_t i = 0; i < input.shape.size() && i < 4; ++i) { - in_ne[i] = input.shape[i]; - } - - int64_t out_ne[4] = {0, 0, 0, 0}; - ref_tensor_f32 out; - if (!sam3_test_run_vit_block_stage(model, block_idx, stage, input.data.data(), in_ne, out.data, out_ne, n_threads)) { - fprintf(stderr, "stage %d for block %d failed\n", (int) stage, block_idx); - std::exit(1); - } - out.shape = trim_shape(out_ne); - return out; -} - -static ref_tensor_f32 run_linear_host_ref_or_die(const sam3_model & model, - int block_idx, - sam3_vit_block_stage stage, - const ref_tensor_f32 & input, - bool use_double_accum) { - int64_t in_ne[4] = {1, 1, 1, 1}; - for (size_t i = 0; i < input.shape.size() && i < 4; ++i) { - in_ne[i] = input.shape[i]; - } - - int64_t out_ne[4] = {0, 0, 0, 0}; - ref_tensor_f32 out; - if (!sam3_test_run_vit_block_linear_host_ref( - model, block_idx, stage, input.data.data(), in_ne, use_double_accum, out.data, out_ne)) { - fprintf(stderr, "host linear ref stage %d for block %d failed\n", (int) stage, block_idx); - std::exit(1); - } - out.shape = trim_shape(out_ne); - return out; -} - -static ref_tensor_f32 add_ref(const ref_tensor_f32 & a, const ref_tensor_f32 & b) { - if (a.shape != b.shape || a.data.size() != b.data.size()) { - fprintf(stderr, "add_ref: shape mismatch\n"); - std::exit(1); - } - ref_tensor_f32 out; - out.shape = a.shape; - out.data.resize(a.data.size()); - for (size_t i = 0; i < a.data.size(); ++i) { - out.data[i] = a.data[i] + b.data[i]; - } - return out; -} - -static ref_tensor_f32 slice_first_columns(const ref_tensor_f32 & t, int n_cols) { - if (t.shape.empty()) { - fprintf(stderr, "slice_first_columns: empty tensor\n"); - std::exit(1); - } - - const int ne0 = t.shape[0]; - const int total_cols = (int) t.data.size() / ne0; - if (total_cols <= 0 || ne0 * total_cols != (int) t.data.size()) { - fprintf(stderr, "slice_first_columns: invalid shape/data size\n"); - std::exit(1); - } - - const int cols = n_cols < total_cols ? n_cols : total_cols; - - ref_tensor_f32 out; - out.shape = {ne0, cols, 1, 1}; - out.data.resize((size_t) ne0 * cols); - std::copy_n(t.data.begin(), (size_t) ne0 * cols, out.data.begin()); - return out; -} - -static ref_tensor_f32 slice_column_range(const ref_tensor_f32 & t, int col0, int n_cols) { - if (t.shape.empty()) { - fprintf(stderr, "slice_column_range: empty tensor\n"); - std::exit(1); - } - - const int ne0 = t.shape[0]; - const int total_cols = (int) t.data.size() / ne0; - if (total_cols <= 0 || ne0 * total_cols != (int) t.data.size()) { - fprintf(stderr, "slice_column_range: invalid shape/data size\n"); - std::exit(1); - } - - if (col0 < 0) { - col0 = 0; - } - if (col0 >= total_cols) { - col0 = total_cols - 1; - } - const int cols = (col0 + n_cols <= total_cols) ? n_cols : (total_cols - col0); - - ref_tensor_f32 out; - out.shape = {ne0, cols, 1, 1}; - out.data.resize((size_t) ne0 * cols); - const float * src = t.data.data() + (size_t) ne0 * col0; - std::copy_n(src, (size_t) ne0 * cols, out.data.begin()); - return out; -} - -static compare_result print_diff_row_ex(const char * label, - const ref_tensor_f32 & a, - const ref_tensor_f32 & b, - float atol = 1e-4f) { - const compare_result r = compare_tensors(a.data.data(), b.data.data(), a.numel(), atol); - fprintf(stderr, "%-26s %-18s %12.6f %12.6f %10d\n", - label, - format_shape(a.shape).c_str(), - r.max_diff, - r.mean_diff, - r.n_bad); - return r; -} - -static void print_diff_row(const char * label, - const ref_tensor_f32 & a, - const ref_tensor_f32 & b, - float atol = 1e-4f) { - (void) print_diff_row_ex(label, a, b, atol); -} - -static void print_graph_row(const dump_item & item, - const std::string & cpu_dir, - const std::string & metal_dir) { - const ref_tensor_f32 cpu = load_ref_f32(cpu_dir + "/" + item.name); - const ref_tensor_f32 metal = load_ref_f32(metal_dir + "/" + item.name); - if (cpu.data.empty() || metal.data.empty()) { - fprintf(stderr, "%-26s %-18s %12s %12s %10s\n", - item.label, "-", "load-fail", "load-fail", "-"); - return; - } - print_diff_row(item.label, cpu, metal); -} - -static bool is_global_block(const sam3_model & model, int block_idx, sam3_tensor_info & rope_info) { - const std::string rope_name = "vit.blocks." + std::to_string(block_idx) + ".attn.freqs_cis"; - if (!sam3_get_model_tensor_info(model, rope_name, rope_info)) { - fprintf(stderr, "failed to load %s info\n", rope_name.c_str()); - std::exit(1); - } - return rope_info.ne[2] == 5184; -} - -struct stage_bundle { - ref_tensor_f32 norm1; - ref_tensor_f32 attn_input; - ref_tensor_f32 qkv; - ref_tensor_f32 attn; - ref_tensor_f32 attn_proj; - ref_tensor_f32 post_attn; - ref_tensor_f32 resid1; - ref_tensor_f32 norm2; - ref_tensor_f32 mlp_fc1; - ref_tensor_f32 mlp_gelu; - ref_tensor_f32 mlp_fc2; - ref_tensor_f32 mlp; - ref_tensor_f32 block; -}; - -static stage_bundle run_block_bundle(const sam3_model & model, - int block_idx, - bool is_global, - const ref_tensor_f32 & block_input, - int n_threads) { - stage_bundle out; - - out.norm1 = run_stage_or_die(model, block_idx, SAM3_VIT_BLOCK_STAGE_NORM1, block_input, n_threads); - out.attn_input = is_global - ? out.norm1 - : run_stage_or_die(model, block_idx, SAM3_VIT_BLOCK_STAGE_WINDOW_PART, out.norm1, n_threads); - out.qkv = run_stage_or_die(model, block_idx, SAM3_VIT_BLOCK_STAGE_QKV_PROJ, out.attn_input, n_threads); - out.attn = run_stage_or_die(model, block_idx, SAM3_VIT_BLOCK_STAGE_ATTN_CORE, out.qkv, n_threads); - out.attn_proj = run_stage_or_die(model, block_idx, SAM3_VIT_BLOCK_STAGE_ATTN_PROJ, out.attn, n_threads); - out.post_attn = is_global - ? out.attn_proj - : run_stage_or_die(model, block_idx, SAM3_VIT_BLOCK_STAGE_WINDOW_UNPART, out.attn_proj, n_threads); - out.resid1 = add_ref(block_input, out.post_attn); - out.norm2 = run_stage_or_die(model, block_idx, SAM3_VIT_BLOCK_STAGE_NORM2, out.resid1, n_threads); - out.mlp_fc1 = run_stage_or_die(model, block_idx, SAM3_VIT_BLOCK_STAGE_MLP_FC1, out.norm2, n_threads); - out.mlp_gelu = run_stage_or_die(model, block_idx, SAM3_VIT_BLOCK_STAGE_MLP_GELU, out.mlp_fc1, n_threads); - out.mlp_fc2 = run_stage_or_die(model, block_idx, SAM3_VIT_BLOCK_STAGE_MLP_FC2, out.mlp_gelu, n_threads); - out.mlp = run_stage_or_die(model, block_idx, SAM3_VIT_BLOCK_STAGE_MLP, out.norm2, n_threads); - out.block = add_ref(out.resid1, out.mlp); - - return out; -} - -int main(int argc, char ** argv) { - if (argc < 2) { - fprintf(stderr, "Usage: %s [ref_dir] [block_idx]\n", argv[0]); - fprintf(stderr, "Default ref_dir: tests/ref_phase3\n"); - fprintf(stderr, "Default block_idx: 16\n"); - return 1; - } - - const std::string model_path = argv[1]; - const std::string ref_dir = argc >= 3 ? argv[2] : "tests/ref_phase3"; - const int block_idx = argc >= 4 ? std::atoi(argv[3]) : 16; - - const std::string input_name = block_idx == 0 ? "dbg_after_ln_pre" : make_block_name(block_idx - 1, "out"); - const std::vector graph_items = { - { "block_input", input_name }, - { "qkv_proj_output", make_block_name(block_idx, "qkv_proj") }, - { "attn_proj_output", make_block_name(block_idx, "attn_proj") }, - { "resid1_output", make_block_name(block_idx, "resid1") }, - { "mlp_output", make_block_name(block_idx, "mlp") }, - { "block_output", make_block_name(block_idx, "out") }, - }; - - const std::string cpu_dir = "/tmp/sam3_block_stage_cpu"; - const std::string metal_dir = "/tmp/sam3_block_stage_metal"; - ensure_dir(cpu_dir); - ensure_dir(metal_dir); - - const ref_tensor_f32 preprocessed = load_ref_f32(ref_dir + "/preprocessed"); - if (preprocessed.data.empty() || preprocessed.shape.size() != 4) { - fprintf(stderr, "failed to load %s/preprocessed\n", ref_dir.c_str()); - return 1; - } - const int img_size = preprocessed.shape[2]; - const int n_threads = 8; - - sam3_params cpu_params; - cpu_params.model_path = model_path; - cpu_params.use_gpu = false; - cpu_params.n_threads = n_threads; - - sam3_params metal_params = cpu_params; - metal_params.use_gpu = true; - - auto cpu_model = sam3_load_model(cpu_params); - auto metal_model = sam3_load_model(metal_params); - if (!cpu_model || !metal_model) { - fprintf(stderr, "failed to load CPU or Metal model\n"); - return 1; - } - - auto cpu_state = sam3_create_state(*cpu_model, cpu_params); - auto metal_state = sam3_create_state(*metal_model, metal_params); - if (!cpu_state || !metal_state) { - fprintf(stderr, "failed to create CPU or Metal state\n"); - return 1; - } - - if (!run_selective_vit(*cpu_state, *cpu_model, preprocessed.data, img_size, block_idx, graph_items, cpu_dir)) { - fprintf(stderr, "CPU selective block run failed\n"); - return 1; - } - if (!run_selective_vit(*metal_state, *metal_model, preprocessed.data, img_size, block_idx, graph_items, metal_dir)) { - fprintf(stderr, "Metal selective block run failed\n"); - return 1; - } - - fprintf(stderr, "\n=== Block-%d graph outputs: CPU vs Metal ===\n", block_idx); - fprintf(stderr, "%-26s %-18s %12s %12s %10s\n", - "checkpoint", "shape", "max_abs_diff", "mean_abs", "n_bad"); - for (const auto & item : graph_items) { - print_graph_row(item, cpu_dir, metal_dir); - } - - const ref_tensor_f32 cpu_block_input = load_ref_f32(cpu_dir + "/" + input_name); - const ref_tensor_f32 metal_block_input = load_ref_f32(metal_dir + "/" + input_name); - if (cpu_block_input.data.empty() || metal_block_input.data.empty()) { - fprintf(stderr, "failed to load canonical block input %s\n", input_name.c_str()); - return 1; - } - - sam3_tensor_info rope_info = {}; - const bool global = is_global_block(*cpu_model, block_idx, rope_info); - - sam3_tensor_info qkv_info = {}; - sam3_tensor_info proj_info = {}; - sam3_tensor_info mlp1_info = {}; - sam3_get_model_tensor_info(*cpu_model, "vit.blocks." + std::to_string(block_idx) + ".attn.qkv.weight", qkv_info); - sam3_get_model_tensor_info(*cpu_model, "vit.blocks." + std::to_string(block_idx) + ".attn.proj.weight", proj_info); - sam3_get_model_tensor_info(*cpu_model, "vit.blocks." + std::to_string(block_idx) + ".mlp.lin1.weight", mlp1_info); - - fprintf(stderr, "\n=== Exact block-%d case ===\n", block_idx); - fprintf(stderr, "block input shape: %s\n", format_shape(cpu_block_input.shape).c_str()); - fprintf(stderr, "attention type: %s\n", global ? "global" : "windowed"); - fprintf(stderr, "qkv weight: shape=[%lld,%lld] type=%d\n", - (long long) qkv_info.ne[0], (long long) qkv_info.ne[1], qkv_info.type); - fprintf(stderr, "proj weight: shape=[%lld,%lld] type=%d\n", - (long long) proj_info.ne[0], (long long) proj_info.ne[1], proj_info.type); - fprintf(stderr, "mlp fc1 weight: shape=[%lld,%lld] type=%d\n", - (long long) mlp1_info.ne[0], (long long) mlp1_info.ne[1], mlp1_info.type); - fprintf(stderr, "freqs_cis: shape=[%lld,%lld,%lld] type=%d\n", - (long long) rope_info.ne[0], (long long) rope_info.ne[1], (long long) rope_info.ne[2], rope_info.type); - print_diff_row("block_input_exact", cpu_block_input, metal_block_input); - - const stage_bundle cpu_shared = run_block_bundle(*cpu_model, block_idx, global, cpu_block_input, n_threads); - const stage_bundle metal_shared = run_block_bundle(*metal_model, block_idx, global, cpu_block_input, n_threads); - const stage_bundle cpu_self = run_block_bundle(*cpu_model, block_idx, global, cpu_block_input, n_threads); - const stage_bundle metal_self = run_block_bundle(*metal_model, block_idx, global, metal_block_input, n_threads); - - fprintf(stderr, "\n=== Exact isolated stages on shared canonical CPU inputs ===\n"); - fprintf(stderr, "%-26s %-18s %12s %12s %10s\n", - "stage", "shape", "max_abs_diff", "mean_abs", "n_bad"); - print_diff_row("norm1_output", cpu_shared.norm1, metal_shared.norm1); - print_diff_row(global ? "attn_input_output" : "window_part_output", cpu_shared.attn_input, metal_shared.attn_input); - const compare_result qkv_diff = print_diff_row_ex("qkv_proj_output", cpu_shared.qkv, metal_shared.qkv); - const compare_result attn_diff = print_diff_row_ex("attn_core_output", cpu_shared.attn, metal_shared.attn); - const compare_result attn_proj_diff = print_diff_row_ex("attn_proj_output", cpu_shared.attn_proj, metal_shared.attn_proj); - print_diff_row(global ? "post_attn_output" : "window_unpart_output", cpu_shared.post_attn, metal_shared.post_attn); - print_diff_row("resid1_output", cpu_shared.resid1, metal_shared.resid1); - print_diff_row("norm2_output", cpu_shared.norm2, metal_shared.norm2); - print_diff_row("mlp_fc1_output", cpu_shared.mlp_fc1, metal_shared.mlp_fc1); - print_diff_row("mlp_gelu_output", cpu_shared.mlp_gelu, metal_shared.mlp_gelu); - const compare_result mlp_fc2_diff = print_diff_row_ex("mlp_fc2_output", cpu_shared.mlp_fc2, metal_shared.mlp_fc2); - print_diff_row("mlp_output", cpu_shared.mlp, metal_shared.mlp); - print_diff_row("block_output", cpu_shared.block, metal_shared.block); - - auto print_worst_loc = [](const char * label, const ref_tensor_f32 & t, const compare_result & diff) { - const int ne0 = t.shape.empty() ? 0 : t.shape[0]; - const int worst_col = ne0 > 0 ? diff.worst_index / ne0 : -1; - const int worst_row = ne0 > 0 ? diff.worst_index % ne0 : -1; - fprintf(stderr, "%s worst: col=%d row=%d cpu=%.6f metal=%.6f max_abs=%.6f\n", - label, worst_col, worst_row, diff.worst_a, diff.worst_b, diff.max_diff); - }; - - fprintf(stderr, "\n=== Worst-element locations on shared canonical CPU inputs ===\n"); - print_worst_loc("qkv_proj_output", cpu_shared.qkv, qkv_diff); - print_worst_loc("attn_core_output", cpu_shared.attn, attn_diff); - print_worst_loc("attn_proj_output", cpu_shared.attn_proj, attn_proj_diff); - print_worst_loc("mlp_fc2_output", cpu_shared.mlp_fc2, mlp_fc2_diff); - - const int linear_slice_cols = 32; - const ref_tensor_f32 qkv_in_slice = slice_first_columns(cpu_shared.attn_input, linear_slice_cols); - const ref_tensor_f32 qkv_cpu_slice = slice_first_columns(cpu_shared.qkv, linear_slice_cols); - const ref_tensor_f32 qkv_mtl_slice = slice_first_columns(metal_shared.qkv, linear_slice_cols); - const ref_tensor_f32 attn_proj_in_slice = slice_first_columns(cpu_shared.attn, linear_slice_cols); - const ref_tensor_f32 attn_proj_cpu_slice = slice_first_columns(cpu_shared.attn_proj, linear_slice_cols); - const ref_tensor_f32 attn_proj_mtl_slice = slice_first_columns(metal_shared.attn_proj, linear_slice_cols); - const ref_tensor_f32 mlp_fc1_in_slice = slice_first_columns(cpu_shared.norm2, linear_slice_cols); - const ref_tensor_f32 mlp_fc1_cpu_slice = slice_first_columns(cpu_shared.mlp_fc1, linear_slice_cols); - const ref_tensor_f32 mlp_fc1_mtl_slice = slice_first_columns(metal_shared.mlp_fc1, linear_slice_cols); - const ref_tensor_f32 mlp_fc2_in_slice = slice_first_columns(cpu_shared.mlp_gelu, linear_slice_cols); - const ref_tensor_f32 mlp_fc2_cpu_slice = slice_first_columns(cpu_shared.mlp_fc2, linear_slice_cols); - const ref_tensor_f32 mlp_fc2_mtl_slice = slice_first_columns(metal_shared.mlp_fc2, linear_slice_cols); - - const ref_tensor_f32 host_qkv_f32 = run_linear_host_ref_or_die(*cpu_model, block_idx, SAM3_VIT_BLOCK_STAGE_QKV_PROJ, qkv_in_slice, false); - const ref_tensor_f32 host_qkv_f64 = run_linear_host_ref_or_die(*cpu_model, block_idx, SAM3_VIT_BLOCK_STAGE_QKV_PROJ, qkv_in_slice, true); - const ref_tensor_f32 host_attn_proj_f32 = run_linear_host_ref_or_die(*cpu_model, block_idx, SAM3_VIT_BLOCK_STAGE_ATTN_PROJ, attn_proj_in_slice, false); - const ref_tensor_f32 host_attn_proj_f64 = run_linear_host_ref_or_die(*cpu_model, block_idx, SAM3_VIT_BLOCK_STAGE_ATTN_PROJ, attn_proj_in_slice, true); - const ref_tensor_f32 host_mlp_fc1_f32 = run_linear_host_ref_or_die(*cpu_model, block_idx, SAM3_VIT_BLOCK_STAGE_MLP_FC1, mlp_fc1_in_slice, false); - const ref_tensor_f32 host_mlp_fc1_f64 = run_linear_host_ref_or_die(*cpu_model, block_idx, SAM3_VIT_BLOCK_STAGE_MLP_FC1, mlp_fc1_in_slice, true); - const ref_tensor_f32 host_mlp_fc2_f32 = run_linear_host_ref_or_die(*cpu_model, block_idx, SAM3_VIT_BLOCK_STAGE_MLP_FC2, mlp_fc2_in_slice, false); - const ref_tensor_f32 host_mlp_fc2_f64 = run_linear_host_ref_or_die(*cpu_model, block_idx, SAM3_VIT_BLOCK_STAGE_MLP_FC2, mlp_fc2_in_slice, true); - - fprintf(stderr, "\n=== Exact block-%d linear host refs (first %d tokens) ===\n", block_idx, linear_slice_cols); - fprintf(stderr, "%-26s %-18s %12s %12s %10s\n", - "stage", "shape", "max_abs_diff", "mean_abs", "n_bad"); - print_diff_row("qkv_cpu_vs_host_f32", qkv_cpu_slice, host_qkv_f32); - print_diff_row("qkv_mtl_vs_host_f32", qkv_mtl_slice, host_qkv_f32); - print_diff_row("qkv_cpu_vs_host_f64", qkv_cpu_slice, host_qkv_f64); - print_diff_row("qkv_mtl_vs_host_f64", qkv_mtl_slice, host_qkv_f64); - print_diff_row("attn_proj_cpu_vs_f32", attn_proj_cpu_slice, host_attn_proj_f32); - print_diff_row("attn_proj_mtl_vs_f32", attn_proj_mtl_slice, host_attn_proj_f32); - print_diff_row("attn_proj_cpu_vs_f64", attn_proj_cpu_slice, host_attn_proj_f64); - print_diff_row("attn_proj_mtl_vs_f64", attn_proj_mtl_slice, host_attn_proj_f64); - print_diff_row("mlp_fc1_cpu_vs_f32", mlp_fc1_cpu_slice, host_mlp_fc1_f32); - print_diff_row("mlp_fc1_mtl_vs_f32", mlp_fc1_mtl_slice, host_mlp_fc1_f32); - print_diff_row("mlp_fc1_cpu_vs_f64", mlp_fc1_cpu_slice, host_mlp_fc1_f64); - print_diff_row("mlp_fc1_mtl_vs_f64", mlp_fc1_mtl_slice, host_mlp_fc1_f64); - print_diff_row("mlp_fc2_cpu_vs_f32", mlp_fc2_cpu_slice, host_mlp_fc2_f32); - print_diff_row("mlp_fc2_mtl_vs_f32", mlp_fc2_mtl_slice, host_mlp_fc2_f32); - print_diff_row("mlp_fc2_cpu_vs_f64", mlp_fc2_cpu_slice, host_mlp_fc2_f64); - print_diff_row("mlp_fc2_mtl_vs_f64", mlp_fc2_mtl_slice, host_mlp_fc2_f64); - - const int qkv_worst_col = qkv_diff.worst_index / cpu_shared.qkv.shape[0]; - const int attn_proj_worst_col = attn_proj_diff.worst_index / cpu_shared.attn_proj.shape[0]; - const int mlp_fc2_worst_col = mlp_fc2_diff.worst_index / cpu_shared.mlp_fc2.shape[0]; - - const ref_tensor_f32 qkv_in_worst = slice_column_range(cpu_shared.attn_input, qkv_worst_col, 1); - const ref_tensor_f32 qkv_cpu_worst = slice_column_range(cpu_shared.qkv, qkv_worst_col, 1); - const ref_tensor_f32 qkv_mtl_worst = slice_column_range(metal_shared.qkv, qkv_worst_col, 1); - const ref_tensor_f32 qkv_host_worst = run_linear_host_ref_or_die(*cpu_model, block_idx, SAM3_VIT_BLOCK_STAGE_QKV_PROJ, qkv_in_worst, false); - - const ref_tensor_f32 attn_proj_in_worst = slice_column_range(cpu_shared.attn, attn_proj_worst_col, 1); - const ref_tensor_f32 attn_proj_cpu_worst = slice_column_range(cpu_shared.attn_proj, attn_proj_worst_col, 1); - const ref_tensor_f32 attn_proj_mtl_worst = slice_column_range(metal_shared.attn_proj, attn_proj_worst_col, 1); - const ref_tensor_f32 attn_proj_host_worst = run_linear_host_ref_or_die(*cpu_model, block_idx, SAM3_VIT_BLOCK_STAGE_ATTN_PROJ, attn_proj_in_worst, false); - - const ref_tensor_f32 mlp_fc2_in_worst = slice_column_range(cpu_shared.mlp_gelu, mlp_fc2_worst_col, 1); - const ref_tensor_f32 mlp_fc2_cpu_worst = slice_column_range(cpu_shared.mlp_fc2, mlp_fc2_worst_col, 1); - const ref_tensor_f32 mlp_fc2_mtl_worst = slice_column_range(metal_shared.mlp_fc2, mlp_fc2_worst_col, 1); - const ref_tensor_f32 mlp_fc2_host_worst = run_linear_host_ref_or_die(*cpu_model, block_idx, SAM3_VIT_BLOCK_STAGE_MLP_FC2, mlp_fc2_in_worst, false); - - fprintf(stderr, "\n=== Worst-column host refs (single real token) ===\n"); - fprintf(stderr, "%-26s %-18s %12s %12s %10s\n", - "stage", "shape", "max_abs_diff", "mean_abs", "n_bad"); - print_diff_row("qkv_cpu_vs_host_worst", qkv_cpu_worst, qkv_host_worst); - print_diff_row("qkv_mtl_vs_host_worst", qkv_mtl_worst, qkv_host_worst); - print_diff_row("attn_proj_cpu_vs_worst", attn_proj_cpu_worst, attn_proj_host_worst); - print_diff_row("attn_proj_mtl_vs_worst", attn_proj_mtl_worst, attn_proj_host_worst); - print_diff_row("mlp_fc2_cpu_vs_worst", mlp_fc2_cpu_worst, mlp_fc2_host_worst); - print_diff_row("mlp_fc2_mtl_vs_worst", mlp_fc2_mtl_worst, mlp_fc2_host_worst); - - const std::map cpu_graph = { - {"qkv_proj_output", load_ref_f32(cpu_dir + "/" + make_block_name(block_idx, "qkv_proj"))}, - {"attn_proj_output", load_ref_f32(cpu_dir + "/" + make_block_name(block_idx, "attn_proj"))}, - {"resid1_output", load_ref_f32(cpu_dir + "/" + make_block_name(block_idx, "resid1"))}, - {"mlp_output", load_ref_f32(cpu_dir + "/" + make_block_name(block_idx, "mlp"))}, - {"block_output", load_ref_f32(cpu_dir + "/" + make_block_name(block_idx, "out"))}, - }; - const std::map metal_graph = { - {"qkv_proj_output", load_ref_f32(metal_dir + "/" + make_block_name(block_idx, "qkv_proj"))}, - {"attn_proj_output", load_ref_f32(metal_dir + "/" + make_block_name(block_idx, "attn_proj"))}, - {"resid1_output", load_ref_f32(metal_dir + "/" + make_block_name(block_idx, "resid1"))}, - {"mlp_output", load_ref_f32(metal_dir + "/" + make_block_name(block_idx, "mlp"))}, - {"block_output", load_ref_f32(metal_dir + "/" + make_block_name(block_idx, "out"))}, - }; - - fprintf(stderr, "\n=== Isolated self-consistency vs graph outputs ===\n"); - fprintf(stderr, "%-26s %-18s %-24s %-24s\n", - "stage", "shape", "cpu_iso_vs_cpu_graph", "metal_iso_vs_metal_graph"); - - auto print_self_row = [&](const char * label, - const ref_tensor_f32 & cpu_iso, - const ref_tensor_f32 & metal_iso, - const char * graph_key) { - const auto cpu_it = cpu_graph.find(graph_key); - const auto metal_it = metal_graph.find(graph_key); - if (cpu_it == cpu_graph.end() || metal_it == metal_graph.end()) { - fprintf(stderr, "%-26s %-18s %-24s %-24s\n", label, "-", "missing", "missing"); - return; - } - - const compare_result cpu_r = compare_tensors(cpu_iso.data.data(), cpu_it->second.data.data(), cpu_iso.numel(), 1e-4f); - const compare_result metal_r = compare_tensors(metal_iso.data.data(), metal_it->second.data.data(), metal_iso.numel(), 1e-4f); - - char cpu_buf[64]; - char metal_buf[64]; - snprintf(cpu_buf, sizeof(cpu_buf), "max=%.6f mean=%.6f", cpu_r.max_diff, cpu_r.mean_diff); - snprintf(metal_buf, sizeof(metal_buf), "max=%.6f mean=%.6f", metal_r.max_diff, metal_r.mean_diff); - fprintf(stderr, "%-26s %-18s %-24s %-24s\n", - label, - format_shape(cpu_iso.shape).c_str(), - cpu_buf, - metal_buf); - }; - - print_self_row("qkv_proj_output", cpu_self.qkv, metal_self.qkv, "qkv_proj_output"); - print_self_row("attn_proj_output", cpu_self.attn_proj, metal_self.attn_proj, "attn_proj_output"); - print_self_row("resid1_output", cpu_self.resid1, metal_self.resid1, "resid1_output"); - print_self_row("mlp_output", cpu_self.mlp, metal_self.mlp, "mlp_output"); - print_self_row("block_output", cpu_self.block, metal_self.block, "block_output"); - - return 0; -} diff --git a/tests/test_metal_conv_dw.cpp b/tests/test_metal_conv_dw.cpp deleted file mode 100644 index 86ab02a..0000000 --- a/tests/test_metal_conv_dw.cpp +++ /dev/null @@ -1,214 +0,0 @@ -// Test Metal kernel for ggml_conv_2d_dw_direct (GGML_OP_CONV_2D_DW). -// Runs the same graph on CPU and Metal, compares results element-by-element. - -#include "ggml-alloc.h" -#include "ggml-backend.h" -#include "ggml-cpu.h" -#include "ggml.h" - -#ifdef GGML_USE_METAL -#include "ggml-metal.h" -#endif - -#include -#include -#include -#include -#include - -static std::vector run_conv_2d_dw(ggml_backend_t backend, - const std::vector& kernel_data, - const std::vector& input_data, - int KW, int KH, int C, - int W, int H, int N, - int s0, int s1, int p0, int p1, - int d0, int d1, - int* out_w, int* out_h) { - size_t ctx_size = ggml_tensor_overhead() * 16 + ggml_graph_overhead(); - struct ggml_init_params params = {ctx_size, nullptr, true}; - struct ggml_context* ctx = ggml_init(params); - - // Kernel: [KW, KH, 1, C] - auto* knl = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, KW, KH, 1, C); - ggml_set_name(knl, "kernel"); - ggml_set_input(knl); - - // Input: [W, H, C, N] - auto* inp = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, W, H, C, N); - ggml_set_name(inp, "input"); - ggml_set_input(inp); - - // Depthwise conv - auto* out = ggml_conv_2d_dw_direct(ctx, knl, inp, s0, s1, p0, p1, d0, d1); - ggml_set_name(out, "output"); - ggml_set_output(out); - - auto* graph = ggml_new_graph(ctx); - ggml_build_forward_expand(graph, out); - - auto* galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!ggml_gallocr_reserve(galloc, graph) || !ggml_gallocr_alloc_graph(galloc, graph)) { - fprintf(stderr, "FAIL: gallocr_alloc_graph failed\n"); - exit(1); - } - - ggml_backend_tensor_set(knl, kernel_data.data(), 0, kernel_data.size() * sizeof(float)); - ggml_backend_tensor_set(inp, input_data.data(), 0, input_data.size() * sizeof(float)); - - ggml_backend_graph_compute(backend, graph); - - *out_w = (int)out->ne[0]; - *out_h = (int)out->ne[1]; - int64_t total = ggml_nelements(out); - - std::vector output(total); - ggml_backend_tensor_get(out, output.data(), 0, total * sizeof(float)); - - ggml_gallocr_free(galloc); - ggml_free(ctx); - return output; -} - -static bool compare_vectors(const std::vector& a, const std::vector& b, - const char* label, float tol = 1e-5f) { - if (a.size() != b.size()) { - fprintf(stderr, " FAIL [%s]: size mismatch: %zu vs %zu\n", label, a.size(), b.size()); - return false; - } - float max_diff = 0.0f; - int first_fail = -1; - for (size_t i = 0; i < a.size(); i++) { - float diff = std::fabs(a[i] - b[i]); - if (diff > max_diff) max_diff = diff; - if (diff > tol && first_fail < 0) first_fail = (int)i; - } - if (first_fail >= 0) { - fprintf(stderr, " FAIL [%s]: max_diff=%.8f at index %d (cpu=%.8f metal=%.8f)\n", - label, max_diff, first_fail, a[first_fail], b[first_fail]); - return false; - } - fprintf(stderr, " PASS [%s]: max_diff=%.2e (%zu elements)\n", label, max_diff, a.size()); - return true; -} - -int main() { -#ifndef GGML_USE_METAL - fprintf(stderr, "Metal not available, skipping test\n"); - return 0; -#else - ggml_backend_t cpu_backend = ggml_backend_cpu_init(); - ggml_backend_t metal_backend = ggml_backend_metal_init(); - if (!metal_backend) { - fprintf(stderr, "FAIL: could not init Metal backend\n"); - return 1; - } - - bool all_pass = true; - - // ── Test 1: sam3 exact dimensions (K=7, C=256, 72x72, s=1, p=3, d=1) ── - { - fprintf(stderr, "\n=== Test 1: sam3 dims (K=7, C=256, 72x72, s=1, p=3) ===\n"); - const int KW = 7, KH = 7, C = 256, W = 72, H = 72, N = 1; - const int s0 = 1, s1 = 1, p0 = 3, p1 = 3, d0 = 1, d1 = 1; - - std::vector kernel_data(KW * KH * 1 * C); - std::vector input_data(W * H * C * N); - for (size_t i = 0; i < kernel_data.size(); i++) - kernel_data[i] = 0.01f * ((float)(i % 97) - 48.0f); - for (size_t i = 0; i < input_data.size(); i++) - input_data[i] = 0.001f * ((float)(i % 9973) - 4986.0f); - - int ow_cpu, oh_cpu, ow_metal, oh_metal; - auto cpu_out = run_conv_2d_dw(cpu_backend, kernel_data, input_data, - KW, KH, C, W, H, N, s0, s1, p0, p1, d0, d1, - &ow_cpu, &oh_cpu); - auto metal_out = run_conv_2d_dw(metal_backend, kernel_data, input_data, - KW, KH, C, W, H, N, s0, s1, p0, p1, d0, d1, - &ow_metal, &oh_metal); - - fprintf(stderr, " output shape: [%d, %d, %d, %d]\n", ow_cpu, oh_cpu, C, N); - all_pass &= compare_vectors(cpu_out, metal_out, "sam3_conv_dw", 1e-4f); - } - - // ── Test 2: small case (K=3, C=4, 8x8, s=1, p=1) ────────────────── - { - fprintf(stderr, "\n=== Test 2: small (K=3, C=4, 8x8, s=1, p=1) ===\n"); - const int KW = 3, KH = 3, C = 4, W = 8, H = 8, N = 1; - const int s0 = 1, s1 = 1, p0 = 1, p1 = 1, d0 = 1, d1 = 1; - - std::vector kernel_data(KW * KH * 1 * C); - std::vector input_data(W * H * C * N); - for (size_t i = 0; i < kernel_data.size(); i++) - kernel_data[i] = (float)(i + 1) * 0.1f; - for (size_t i = 0; i < input_data.size(); i++) - input_data[i] = (float)(i + 1) * 0.01f; - - int ow_cpu, oh_cpu, ow_metal, oh_metal; - auto cpu_out = run_conv_2d_dw(cpu_backend, kernel_data, input_data, - KW, KH, C, W, H, N, s0, s1, p0, p1, d0, d1, - &ow_cpu, &oh_cpu); - auto metal_out = run_conv_2d_dw(metal_backend, kernel_data, input_data, - KW, KH, C, W, H, N, s0, s1, p0, p1, d0, d1, - &ow_metal, &oh_metal); - - fprintf(stderr, " output shape: [%d, %d, %d, %d]\n", ow_cpu, oh_cpu, C, N); - all_pass &= compare_vectors(cpu_out, metal_out, "small_conv_dw", 1e-5f); - } - - // ── Test 3: stride=2, dilation=2 ──────────────────────────────────── - { - fprintf(stderr, "\n=== Test 3: stride=2, dilation=2 (K=3, C=16, 32x32) ===\n"); - const int KW = 3, KH = 3, C = 16, W = 32, H = 32, N = 1; - const int s0 = 2, s1 = 2, p0 = 2, p1 = 2, d0 = 2, d1 = 2; - - std::vector kernel_data(KW * KH * 1 * C); - std::vector input_data(W * H * C * N); - for (size_t i = 0; i < kernel_data.size(); i++) - kernel_data[i] = sinf((float)i * 0.1f); - for (size_t i = 0; i < input_data.size(); i++) - input_data[i] = cosf((float)i * 0.01f); - - int ow_cpu, oh_cpu, ow_metal, oh_metal; - auto cpu_out = run_conv_2d_dw(cpu_backend, kernel_data, input_data, - KW, KH, C, W, H, N, s0, s1, p0, p1, d0, d1, - &ow_cpu, &oh_cpu); - auto metal_out = run_conv_2d_dw(metal_backend, kernel_data, input_data, - KW, KH, C, W, H, N, s0, s1, p0, p1, d0, d1, - &ow_metal, &oh_metal); - - fprintf(stderr, " output shape: [%d, %d, %d, %d]\n", ow_cpu, oh_cpu, C, N); - all_pass &= compare_vectors(cpu_out, metal_out, "strided_dilated", 1e-5f); - } - - // ── Test 4: no padding ────────────────────────────────────────────── - { - fprintf(stderr, "\n=== Test 4: no padding (K=5, C=8, 16x16) ===\n"); - const int KW = 5, KH = 5, C = 8, W = 16, H = 16, N = 2; - const int s0 = 1, s1 = 1, p0 = 0, p1 = 0, d0 = 1, d1 = 1; - - std::vector kernel_data(KW * KH * 1 * C); - std::vector input_data(W * H * C * N); - for (size_t i = 0; i < kernel_data.size(); i++) - kernel_data[i] = 0.02f * ((float)(i % 51) - 25.0f); - for (size_t i = 0; i < input_data.size(); i++) - input_data[i] = 0.005f * ((float)(i % 199) - 99.0f); - - int ow_cpu, oh_cpu, ow_metal, oh_metal; - auto cpu_out = run_conv_2d_dw(cpu_backend, kernel_data, input_data, - KW, KH, C, W, H, N, s0, s1, p0, p1, d0, d1, - &ow_cpu, &oh_cpu); - auto metal_out = run_conv_2d_dw(metal_backend, kernel_data, input_data, - KW, KH, C, W, H, N, s0, s1, p0, p1, d0, d1, - &ow_metal, &oh_metal); - - fprintf(stderr, " output shape: [%d, %d, %d, %d]\n", ow_cpu, oh_cpu, C, N); - all_pass &= compare_vectors(cpu_out, metal_out, "no_pad_batch2", 1e-5f); - } - - ggml_backend_free(cpu_backend); - ggml_backend_free(metal_backend); - - fprintf(stderr, "\n%s\n", all_pass ? "ALL TESTS PASSED" : "SOME TESTS FAILED"); - return all_pass ? 0 : 1; -#endif -} diff --git a/tests/test_metal_conv_transpose_2d.cpp b/tests/test_metal_conv_transpose_2d.cpp deleted file mode 100644 index 31b9838..0000000 --- a/tests/test_metal_conv_transpose_2d.cpp +++ /dev/null @@ -1,338 +0,0 @@ -#include "ggml-alloc.h" -#include "ggml-backend.h" -#include "ggml-cpu.h" -#include "ggml.h" - -#ifdef GGML_USE_METAL -#include "ggml-metal.h" -#endif - -#include "test_utils.h" - -#include -#include -#include -#include -#include -#include -#include -#include - -struct options { - std::string case_name = "all"; - bool dump_mismatch = false; - int n_threads = 8; - int max_wall_ms = 30000; -}; - -struct case_desc { - const char * name; - int64_t in_w; - int64_t in_h; - int64_t in_c; - int64_t out_c; - int64_t kw; - int64_t kh; - int stride; - uint32_t seed; - float tol; -}; - -struct run_result { - std::vector output; - int64_t out_w = 0; - int64_t out_h = 0; - int64_t out_c = 0; -}; - -static const case_desc CASES[] = { - { "cpu_ref_stride1", 3, 2, 2, 3, 2, 2, 1, 0x1234u, 1e-5f }, - { "cpu_ref_stride2", 3, 2, 2, 3, 2, 2, 2, 0x2345u, 1e-5f }, - { "cpu_ref_stride3", 3, 2, 2, 3, 2, 2, 3, 0x3456u, 1e-5f }, - { "small_rect", 5, 4, 3, 4, 3, 2, 1, 0x4567u, 1e-5f }, - { "medium_stride2", 8, 7, 5, 6, 3, 3, 2, 0x5678u, 1e-4f }, - { "sam_first", 72, 72, 1024, 512, 2, 2, 2, 0x6789u, 2e-3f }, - { "sam_second", 144, 144, 512, 256, 2, 2, 2, 0x789au, 2e-3f }, -}; - -static void usage(const char * argv0) { - fprintf(stderr, - "Usage: %s [--case NAME|all] [--dump-mismatch] [--n-threads N] [--max-wall-ms N] [--list]\n", - argv0); -} - -static bool parse_args(int argc, char ** argv, options & opts) { - for (int i = 1; i < argc; ++i) { - if (strcmp(argv[i], "--case") == 0) { - if (i + 1 >= argc) return false; - opts.case_name = argv[++i]; - } else if (strcmp(argv[i], "--dump-mismatch") == 0) { - opts.dump_mismatch = true; - } else if (strcmp(argv[i], "--n-threads") == 0) { - if (i + 1 >= argc) return false; - opts.n_threads = atoi(argv[++i]); - } else if (strcmp(argv[i], "--max-wall-ms") == 0) { - if (i + 1 >= argc) return false; - opts.max_wall_ms = atoi(argv[++i]); - } else if (strcmp(argv[i], "--list") == 0) { - for (const case_desc & tc : CASES) { - fprintf(stderr, "%s\n", tc.name); - } - exit(0); - } else if (strcmp(argv[i], "--help") == 0 || strcmp(argv[i], "-h") == 0) { - return false; - } else { - fprintf(stderr, "Unknown argument: %s\n", argv[i]); - return false; - } - } - return true; -} - -static float gen_value(uint64_t idx, uint32_t seed, float scale) { - uint32_t x = (uint32_t) idx ^ seed; - x = x * 1664525u + 1013904223u; - x ^= x >> 16; - const int32_t centered = (int32_t) (x % 511u) - 255; - return centered * scale; -} - -static void fill_input_data(const case_desc & tc, std::vector & input) { - for (size_t i = 0; i < input.size(); ++i) { - input[i] = gen_value(i, tc.seed ^ 0x13579bdu, 1.0f / 256.0f); - } -} - -static void fill_weight_data(const case_desc & tc, std::vector & weight_f16, std::vector & weight_f32) { - for (size_t i = 0; i < weight_f16.size(); ++i) { - const float v = gen_value(i, tc.seed ^ 0x2468aceu, 1.0f / 128.0f); - weight_f32[i] = v; - weight_f16[i] = ggml_fp32_to_fp16(v); - } -} - -static ggml_backend_t create_backend(const std::string & name, int n_threads) { - if (name == "cpu") { - ggml_backend_t backend = ggml_backend_cpu_init(); - if (backend) { - ggml_backend_cpu_set_n_threads(backend, n_threads); - } - return backend; - } - -#ifdef GGML_USE_METAL - if (name == "metal") { - return ggml_backend_metal_init(); - } -#endif - - return nullptr; -} - -static run_result run_case(const case_desc & tc, - const std::string & backend_name, - int n_threads, - const std::vector & weight_data, - const std::vector & input_data) { - ggml_backend_t backend = create_backend(backend_name, n_threads); - if (!backend) { - fprintf(stderr, "Failed to create backend '%s'\n", backend_name.c_str()); - exit(2); - } - - const size_t ctx_size = ggml_tensor_overhead() * 8 + ggml_graph_overhead(); - ggml_init_params params = { - /*.mem_size =*/ ctx_size, - /*.mem_buffer =*/ nullptr, - /*.no_alloc =*/ true, - }; - - ggml_context * ctx = ggml_init(params); - if (!ctx) { - fprintf(stderr, "Failed to create ggml context\n"); - ggml_backend_free(backend); - exit(3); - } - - ggml_tensor * weight = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, tc.kw, tc.kh, tc.out_c, tc.in_c); - ggml_tensor * input = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, tc.in_w, tc.in_h, tc.in_c, 1); - ggml_set_name(weight, "weight"); - ggml_set_name(input, "input"); - ggml_set_input(weight); - ggml_set_input(input); - - ggml_tensor * output = ggml_conv_transpose_2d_p0(ctx, weight, input, tc.stride); - ggml_set_name(output, "output"); - ggml_set_output(output); - - ggml_cgraph * graph = ggml_new_graph(ctx); - ggml_build_forward_expand(graph, output); - - ggml_gallocr_t galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!ggml_gallocr_reserve(galloc, graph) || !ggml_gallocr_alloc_graph(galloc, graph)) { - fprintf(stderr, "Failed to allocate graph for case %s (%s)\n", tc.name, backend_name.c_str()); - ggml_gallocr_free(galloc); - ggml_free(ctx); - ggml_backend_free(backend); - exit(4); - } - - ggml_backend_tensor_set(weight, weight_data.data(), 0, weight_data.size() * sizeof(uint16_t)); - ggml_backend_tensor_set(input, input_data.data(), 0, input_data.size() * sizeof(float)); - - const ggml_status status = ggml_backend_graph_compute(backend, graph); - if (status != GGML_STATUS_SUCCESS) { - fprintf(stderr, "Graph compute failed for case %s on %s: %s\n", - tc.name, backend_name.c_str(), ggml_status_to_string(status)); - ggml_gallocr_free(galloc); - ggml_free(ctx); - ggml_backend_free(backend); - exit(5); - } - - run_result res; - res.out_w = output->ne[0]; - res.out_h = output->ne[1]; - res.out_c = output->ne[2]; - res.output.resize((size_t) ggml_nelements(output)); - ggml_backend_tensor_get(output, res.output.data(), 0, res.output.size() * sizeof(float)); - - ggml_gallocr_free(galloc); - ggml_free(ctx); - ggml_backend_free(backend); - return res; -} - -static void dump_reference_mismatch(const case_desc & tc, - const std::vector & input_data, - const std::vector & weight_data, - const run_result & cpu, - const run_result & metal, - int worst_index) { - if (worst_index < 0) { - return; - } - - const int64_t ow = worst_index % cpu.out_w; - const int64_t rem0 = worst_index / cpu.out_w; - const int64_t oh = rem0 % cpu.out_h; - const int64_t oc = rem0 / cpu.out_h; - - fprintf(stderr, - " mismatch case=%s coord=(ow=%lld, oh=%lld, oc=%lld) cpu=%.8f metal=%.8f\n", - tc.name, - (long long) ow, - (long long) oh, - (long long) oc, - cpu.output[worst_index], - metal.output[worst_index]); - - int shown = 0; - double recomputed = 0.0; - for (int64_t kh = 0; kh < tc.kh; ++kh) { - const int64_t in_y_nom = oh - kh; - if (in_y_nom < 0 || in_y_nom % tc.stride != 0) { - continue; - } - const int64_t in_y = in_y_nom / tc.stride; - if (in_y >= tc.in_h) { - continue; - } - - for (int64_t kw = 0; kw < tc.kw; ++kw) { - const int64_t in_x_nom = ow - kw; - if (in_x_nom < 0 || in_x_nom % tc.stride != 0) { - continue; - } - const int64_t in_x = in_x_nom / tc.stride; - if (in_x >= tc.in_w) { - continue; - } - - for (int64_t ic = 0; ic < tc.in_c; ++ic) { - const size_t input_idx = (size_t) ic * tc.in_h * tc.in_w + (size_t) in_y * tc.in_w + (size_t) in_x; - const size_t weight_idx = ((size_t) ic * tc.out_c * tc.kh * tc.kw) + - ((size_t) oc * tc.kh * tc.kw) + - (size_t) kh * tc.kw + - (size_t) kw; - const float contrib = input_data[input_idx] * weight_data[weight_idx]; - recomputed += contrib; - if (shown < 16) { - fprintf(stderr, - " term[%2d] in=(x=%lld,y=%lld,c=%lld)=%.8f weight=(kw=%lld,kh=%lld,oc=%lld,ic=%lld)=%.8f contrib=%.8f\n", - shown, - (long long) in_x, - (long long) in_y, - (long long) ic, - input_data[input_idx], - (long long) kw, - (long long) kh, - (long long) oc, - (long long) ic, - weight_data[weight_idx], - contrib); - } - ++shown; - } - } - } - - fprintf(stderr, " recomputed_reference=%.8f from %d contributing terms\n", recomputed, shown); -} - -int main(int argc, char ** argv) { - options opts; - if (!parse_args(argc, argv, opts)) { - usage(argv[0]); - return 1; - } - -#ifndef GGML_USE_METAL - fprintf(stderr, "Metal not available, skipping test\n"); - return 0; -#else - bool all_pass = true; - - fprintf(stderr, "| case | output | max abs diff | mean abs diff | tolerance | status |\n"); - fprintf(stderr, "| --- | --- | ---: | ---: | ---: | --- |\n"); - - for (const case_desc & tc : CASES) { - if (opts.case_name != "all" && opts.case_name != tc.name) { - continue; - } - - std::vector weight_f16((size_t) tc.kw * tc.kh * tc.out_c * tc.in_c); - std::vector weight_f32(weight_f16.size()); - std::vector input_data((size_t) tc.in_w * tc.in_h * tc.in_c); - - fill_weight_data(tc, weight_f16, weight_f32); - fill_input_data(tc, input_data); - - run_result cpu = run_case(tc, "cpu", opts.n_threads, weight_f16, input_data); - run_result metal = run_case(tc, "metal", opts.n_threads, weight_f16, input_data); - - compare_result diff = compare_tensors(cpu.output.data(), metal.output.data(), (int) cpu.output.size(), tc.tol); - const bool pass = diff.max_diff <= tc.tol; - fprintf(stderr, - "| %s | [%lld,%lld,%lld] | %.8f | %.8f | %.1e | %s |\n", - tc.name, - (long long) cpu.out_w, - (long long) cpu.out_h, - (long long) cpu.out_c, - diff.max_diff, - diff.mean_diff, - tc.tol, - pass ? "PASS" : "FAIL"); - - if (!pass) { - all_pass = false; - if (opts.dump_mismatch || (tc.in_w <= 8 && tc.in_h <= 8 && tc.in_c <= 8)) { - dump_reference_mismatch(tc, input_data, weight_f32, cpu, metal, diff.worst_index); - } - } - } - - fprintf(stderr, "\n%s\n", all_pass ? "ALL TESTS PASSED" : "SOME TESTS FAILED"); - return all_pass ? 0 : 1; -#endif -} diff --git a/tests/test_metal_conv_transpose_2d_stress.cpp b/tests/test_metal_conv_transpose_2d_stress.cpp deleted file mode 100644 index 76f1207..0000000 --- a/tests/test_metal_conv_transpose_2d_stress.cpp +++ /dev/null @@ -1,219 +0,0 @@ -#include "ggml-alloc.h" -#include "ggml-backend.h" -#include "ggml-cpu.h" -#include "ggml.h" - -#ifdef GGML_USE_METAL -#include "ggml-metal.h" -#endif - -#include -#include -#include -#include -#include -#include - -struct options { - std::string backend = "metal"; - std::string variant = "first"; - int warmup = 1; - int repeats = 3; - int n_threads = 8; - int max_wall_ms = 0; -}; - -struct variant_desc { - const char * name; - int64_t in_w; - int64_t in_h; - int64_t in_c; - int64_t out_c; -}; - -static const variant_desc VARIANT_FIRST = { "first", 72, 72, 1024, 512 }; -static const variant_desc VARIANT_SECOND = { "second", 144, 144, 512, 256 }; - -static void usage(const char * argv0) { - fprintf(stderr, - "Usage: %s [--backend metal|cpu] [--variant first|second] [--warmup N] [--repeats N] [--max-wall-ms N]\n", - argv0); -} - -static bool parse_args(int argc, char ** argv, options & opts) { - for (int i = 1; i < argc; ++i) { - if (strcmp(argv[i], "--backend") == 0) { - if (i + 1 >= argc) return false; - opts.backend = argv[++i]; - } else if (strcmp(argv[i], "--variant") == 0) { - if (i + 1 >= argc) return false; - opts.variant = argv[++i]; - } else if (strcmp(argv[i], "--warmup") == 0) { - if (i + 1 >= argc) return false; - opts.warmup = atoi(argv[++i]); - } else if (strcmp(argv[i], "--repeats") == 0) { - if (i + 1 >= argc) return false; - opts.repeats = atoi(argv[++i]); - } else if (strcmp(argv[i], "--n-threads") == 0) { - if (i + 1 >= argc) return false; - opts.n_threads = atoi(argv[++i]); - } else if (strcmp(argv[i], "--max-wall-ms") == 0) { - if (i + 1 >= argc) return false; - opts.max_wall_ms = atoi(argv[++i]); - } else if (strcmp(argv[i], "--help") == 0 || strcmp(argv[i], "-h") == 0) { - return false; - } else { - fprintf(stderr, "Unknown argument: %s\n", argv[i]); - return false; - } - } - - if (opts.backend != "metal" && opts.backend != "cpu") { - return false; - } - - if (opts.variant != "first" && opts.variant != "second") { - return false; - } - - return true; -} - -static const variant_desc & get_variant(const std::string & name) { - return name == "second" ? VARIANT_SECOND : VARIANT_FIRST; -} - -static ggml_backend_t create_backend(const options & opts) { -#ifdef GGML_USE_METAL - if (opts.backend == "metal") { - return ggml_backend_metal_init(); - } -#endif - return ggml_backend_cpu_init(); -} - -int main(int argc, char ** argv) { - options opts; - if (!parse_args(argc, argv, opts)) { - usage(argv[0]); - return 1; - } - - ggml_backend_t backend = create_backend(opts); - if (!backend) { - fprintf(stderr, "Failed to create backend '%s'\n", opts.backend.c_str()); - return 2; - } - - if (ggml_backend_is_cpu(backend)) { - ggml_backend_cpu_set_n_threads(backend, opts.n_threads); - } - - const variant_desc & v = get_variant(opts.variant); - const int64_t kw = 2; - const int64_t kh = 2; - const int stride = 2; - - const size_t ctx_size = ggml_tensor_overhead() * 8 + ggml_graph_overhead(); - ggml_init_params params = { - /*.mem_size =*/ ctx_size, - /*.mem_buffer =*/ nullptr, - /*.no_alloc =*/ true, - }; - ggml_context * ctx = ggml_init(params); - if (!ctx) { - fprintf(stderr, "Failed to create ggml context\n"); - ggml_backend_free(backend); - return 3; - } - - ggml_tensor * weight = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, kw, kh, v.out_c, v.in_c); - ggml_tensor * input = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, v.in_w, v.in_h, v.in_c, 1); - ggml_set_name(weight, "weight"); - ggml_set_name(input, "input"); - ggml_set_input(weight); - ggml_set_input(input); - - ggml_tensor * output = ggml_conv_transpose_2d_p0(ctx, weight, input, stride); - ggml_set_name(output, "output"); - ggml_set_output(output); - - ggml_cgraph * graph = ggml_new_graph(ctx); - ggml_build_forward_expand(graph, output); - - ggml_gallocr_t galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!ggml_gallocr_reserve(galloc, graph) || !ggml_gallocr_alloc_graph(galloc, graph)) { - fprintf(stderr, "Failed to allocate graph\n"); - ggml_gallocr_free(galloc); - ggml_free(ctx); - ggml_backend_free(backend); - return 4; - } - - std::vector wdata((size_t) ggml_nelements(weight)); - for (size_t i = 0; i < wdata.size(); ++i) { - const int32_t centered = (int32_t) (i % 17) - 8; - wdata[i] = ggml_fp32_to_fp16((float) centered * 0.015625f); - } - std::vector xdata((size_t) ggml_nelements(input)); - for (size_t i = 0; i < xdata.size(); ++i) { - const int32_t centered = (int32_t) (i % 23) - 11; - xdata[i] = (float) centered * 0.0078125f; - } - - ggml_backend_tensor_set(weight, wdata.data(), 0, wdata.size() * sizeof(uint16_t)); - ggml_backend_tensor_set(input, xdata.data(), 0, xdata.size() * sizeof(float)); - - for (int i = 0; i < opts.warmup; ++i) { - const ggml_status status = ggml_backend_graph_compute(backend, graph); - if (status != GGML_STATUS_SUCCESS) { - fprintf(stderr, "Warmup failed: %s\n", ggml_status_to_string(status)); - ggml_gallocr_free(galloc); - ggml_free(ctx); - ggml_backend_free(backend); - return 5; - } - } - - auto t0 = std::chrono::high_resolution_clock::now(); - for (int i = 0; i < opts.repeats; ++i) { - const ggml_status status = ggml_backend_graph_compute(backend, graph); - if (status != GGML_STATUS_SUCCESS) { - fprintf(stderr, "Run %d failed: %s\n", i + 1, ggml_status_to_string(status)); - ggml_gallocr_free(galloc); - ggml_free(ctx); - ggml_backend_free(backend); - return 6; - } - } - auto t1 = std::chrono::high_resolution_clock::now(); - - std::vector out((size_t) std::min(ggml_nelements(output), 64)); - ggml_backend_tensor_get(output, out.data(), 0, out.size() * sizeof(float)); - double checksum = 0.0; - for (float vout : out) { - checksum += vout; - } - - const double avg_ms = std::chrono::duration(t1 - t0).count() / opts.repeats; - - fprintf(stderr, - "RESULT backend=%s variant=%s in=[%lld,%lld,%lld] out_c=%lld out=[%lld,%lld,%lld] avg_ms=%.3f checksum=%.6f repeats=%d\n", - opts.backend.c_str(), - v.name, - (long long) v.in_w, - (long long) v.in_h, - (long long) v.in_c, - (long long) v.out_c, - (long long) output->ne[0], - (long long) output->ne[1], - (long long) output->ne[2], - avg_ms, - checksum, - opts.repeats); - - ggml_gallocr_free(galloc); - ggml_free(ctx); - ggml_backend_free(backend); - return 0; -} diff --git a/tests/test_metal_encoder_checkpoints.cpp b/tests/test_metal_encoder_checkpoints.cpp deleted file mode 100644 index cbbd0be..0000000 --- a/tests/test_metal_encoder_checkpoints.cpp +++ /dev/null @@ -1,203 +0,0 @@ -#include "sam3.h" -#include "test_utils.h" - -#include -#include -#include -#include -#include - -struct checkpoint_stat { - std::string name; - std::string shape; - compare_result diff; -}; - -static std::string format_shape(const std::vector & shape) { - std::string out = "["; - for (size_t i = 0; i < shape.size(); ++i) { - if (i > 0) { - out += ", "; - } - out += std::to_string(shape[i]); - } - out += "]"; - return out; -} - -static bool dump_checkpoints(const sam3_state & state, - const std::vector & checkpoints, - const std::string & output_dir) { - for (const auto & name : checkpoints) { - if (!sam3_dump_state_tensor(state, name, output_dir + "/" + name)) { - fprintf(stderr, "failed to dump checkpoint '%s'\n", name.c_str()); - return false; - } - } - - return true; -} - -static bool run_backend(const std::string & model_path, - const std::vector & chw_data, - int img_size, - bool use_gpu, - int n_threads, - const std::vector & checkpoints, - const std::string & dump_dir, - double & elapsed_ms) { - sam3_params params; - params.model_path = model_path; - params.use_gpu = use_gpu; - params.n_threads = n_threads; - - auto model = sam3_load_model(params); - if (!model) { - fprintf(stderr, "failed to load model for %s\n", use_gpu ? "Metal" : "CPU"); - return false; - } - - auto state = sam3_create_state(*model, params); - if (!state) { - fprintf(stderr, "failed to create state for %s\n", use_gpu ? "Metal" : "CPU"); - return false; - } - - auto t0 = std::chrono::high_resolution_clock::now(); - const bool ok = sam3_encode_image_from_preprocessed(*state, *model, chw_data.data(), img_size); - auto t1 = std::chrono::high_resolution_clock::now(); - elapsed_ms = std::chrono::duration(t1 - t0).count(); - if (!ok) { - fprintf(stderr, "%s encode failed\n", use_gpu ? "Metal" : "CPU"); - return false; - } - - if (!dump_checkpoints(*state, checkpoints, dump_dir)) { - return false; - } - - state.reset(); - sam3_free_model(*model); - model.reset(); - return true; -} - -int main(int argc, char ** argv) { - if (argc < 2) { - fprintf(stderr, "Usage: %s [ref_dir]\n", argv[0]); - fprintf(stderr, "Default ref_dir: tests/ref_phase3\n"); - return 1; - } - - const std::string model_path = argv[1]; - const std::string ref_dir = argc >= 3 ? argv[2] : "tests/ref_phase3"; - - auto preprocessed = load_ref_f32(ref_dir + "/preprocessed"); - if (preprocessed.data.empty()) { - fprintf(stderr, "failed to load %s/preprocessed.bin\n", ref_dir.c_str()); - return 1; - } - if (preprocessed.shape.size() != 4) { - fprintf(stderr, "unexpected preprocessed shape rank: %zu\n", preprocessed.shape.size()); - return 1; - } - - const int img_size = preprocessed.shape[2]; - std::vector checkpoints = { - "dbg_patch_embed", - "dbg_after_pos_embed", - "dbg_ln_pre_norm", - "dbg_ln_pre_scale", - "dbg_after_ln_pre", - "dbg_block_15_norm1", - "dbg_block_15_attn_out", - "dbg_block_15_attn_proj", - "dbg_block_15_resid1", - "dbg_block_15_norm2", - "dbg_block_15_mlp", - }; - for (int i = 0; i < 31; ++i) { - checkpoints.emplace_back("dbg_block_" + std::to_string(i) + "_out"); - } - checkpoints.emplace_back("vit_output"); - checkpoints.emplace_back("neck_trk_0"); - checkpoints.emplace_back("neck_trk_1"); - checkpoints.emplace_back("neck_trk_2"); - - const std::string cpu_dir = "/tmp/sam3_cpu_checkpoints"; - const std::string metal_dir = "/tmp/sam3_metal_checkpoints"; - ensure_dir(cpu_dir); - ensure_dir(metal_dir); - - double cpu_ms = 0.0; - double metal_ms = 0.0; - - fprintf(stderr, "\n=== CPU Run ===\n"); - if (!run_backend(model_path, preprocessed.data, img_size, false, 8, checkpoints, cpu_dir, cpu_ms)) { - return 1; - } - fprintf(stderr, "CPU encoder wall time: %.1f ms\n", cpu_ms); - - fprintf(stderr, "\n=== Metal Run ===\n"); - if (!run_backend(model_path, preprocessed.data, img_size, true, 8, checkpoints, metal_dir, metal_ms)) { - return 1; - } - fprintf(stderr, "Metal encoder wall time: %.1f ms\n", metal_ms); - - fprintf(stderr, "\n=== CPU vs Metal Checkpoints ===\n"); - fprintf(stderr, "%-22s %-18s %14s %14s %12s\n", - "checkpoint", "shape", "max_abs_diff", "mean_abs_diff", "n_bad"); - - std::vector stats; - stats.reserve(checkpoints.size()); - - for (const auto & name : checkpoints) { - auto cpu = load_ref_f32(cpu_dir + "/" + name); - auto metal = load_ref_f32(metal_dir + "/" + name); - if (cpu.data.empty() || metal.data.empty()) { - fprintf(stderr, "%-22s %-18s %14s %14s %12s\n", - name.c_str(), "-", "load-fail", "load-fail", "-"); - continue; - } - if (cpu.shape != metal.shape) { - fprintf(stderr, "%-22s %-18s %14s %14s %12s\n", - name.c_str(), "shape-mismatch", "shape-mismatch", "shape-mismatch", "-"); - continue; - } - - checkpoint_stat stat; - stat.name = name; - stat.shape = format_shape(cpu.shape); - stat.diff = compare_tensors(cpu.data.data(), metal.data.data(), cpu.numel(), 1e-4f); - stats.push_back(stat); - - fprintf(stderr, "%-22s %-18s %14.6f %14.6f %12d\n", - stat.name.c_str(), - stat.shape.c_str(), - stat.diff.max_diff, - stat.diff.mean_diff, - stat.diff.n_bad); - } - - const float checkpoint_tol = 1e-2f; - const checkpoint_stat * first_mismatch = nullptr; - for (const auto & stat : stats) { - if (stat.diff.max_diff > checkpoint_tol) { - first_mismatch = &stat; - break; - } - } - - if (first_mismatch) { - fprintf(stderr, - "\nFIRST CLEAR MISMATCH: %s shape=%s max_abs_diff=%.6f mean_abs_diff=%.6f\n", - first_mismatch->name.c_str(), - first_mismatch->shape.c_str(), - first_mismatch->diff.max_diff, - first_mismatch->diff.mean_diff); - return 1; - } - - fprintf(stderr, "\nAll checkpoints stayed within %.3g max abs diff.\n", checkpoint_tol); - return 0; -} diff --git a/tests/test_metal_flash_attn_ext_block15_compare.cpp b/tests/test_metal_flash_attn_ext_block15_compare.cpp deleted file mode 100644 index 2e90fae..0000000 --- a/tests/test_metal_flash_attn_ext_block15_compare.cpp +++ /dev/null @@ -1,715 +0,0 @@ -#include "ggml-alloc.h" -#include "ggml-backend.h" -#include "ggml-cpu.h" -#include "ggml.h" - -#ifdef GGML_USE_METAL -#include "ggml-metal.h" -#endif - -#include "sam3.h" -#include "test_utils.h" - -#include -#include -#include -#include -#include -#include -#include - -struct tensor_case { - ref_tensor_f32 qkv_proj; - ref_tensor_f32 freqs_cis; - ref_tensor_f32 q_heads; - ref_tensor_f32 k_heads; - ref_tensor_f32 q_rope; - ref_tensor_f32 k_rope; - ref_tensor_f32 q_flash; - ref_tensor_f32 k_flash; - ref_tensor_f32 v_base; - ref_tensor_f32 v_flash; - ref_tensor_f32 attn_out; -}; - -struct prep_run_result { - std::map tensors; -}; - -struct backend_run_result { - double elapsed_ms = 0.0; - std::map infos; -}; - -struct isolated_run_result { - double elapsed_ms = 0.0; - std::vector output; - std::vector shape; -}; - -static void usage(const char * argv0) { - fprintf(stderr, "Usage: %s [ref_dir]\n", argv0); - fprintf(stderr, "Default ref_dir: tests/ref_phase3\n"); -} - -static std::string format_shape(const std::vector & shape) { - std::string out = "["; - for (size_t i = 0; i < shape.size(); ++i) { - if (i > 0) { - out += ", "; - } - out += std::to_string(shape[i]); - } - out += "]"; - return out; -} - -static const char * ggml_type_name_public(int type) { - switch ((ggml_type) type) { - case GGML_TYPE_F32: return "F32"; - case GGML_TYPE_F16: return "F16"; - case GGML_TYPE_BF16: return "BF16"; - case GGML_TYPE_I32: return "I32"; - default: return "other"; - } -} - -static bool dump_named_tensors(const sam3_state & state, - const std::vector & names, - const std::string & output_dir, - std::map & infos) { - for (const auto & name : names) { - sam3_tensor_info info; - if (!sam3_get_state_tensor_info(state, name, info)) { - fprintf(stderr, "missing tensor info for %s\n", name.c_str()); - return false; - } - infos[name] = info; - - if (!sam3_dump_state_tensor(state, name, output_dir + "/" + name)) { - fprintf(stderr, "failed to dump %s\n", name.c_str()); - return false; - } - } - return true; -} - -static bool run_encoder_backend(const std::string & model_path, - const std::vector & chw_data, - int img_size, - bool use_gpu, - int n_threads, - const std::vector & names, - const std::string & dump_dir, - backend_run_result & result) { - sam3_params params; - params.model_path = model_path; - params.use_gpu = use_gpu; - params.n_threads = n_threads; - - auto model = sam3_load_model(params); - if (!model) { - fprintf(stderr, "failed to load model for %s\n", use_gpu ? "Metal" : "CPU"); - return false; - } - - auto state = sam3_create_state(*model, params); - if (!state) { - fprintf(stderr, "failed to create state for %s\n", use_gpu ? "Metal" : "CPU"); - return false; - } - - auto t0 = std::chrono::high_resolution_clock::now(); - const bool ok = sam3_encode_image_from_preprocessed(*state, *model, chw_data.data(), img_size); - auto t1 = std::chrono::high_resolution_clock::now(); - result.elapsed_ms = std::chrono::duration(t1 - t0).count(); - if (!ok) { - fprintf(stderr, "%s encode failed\n", use_gpu ? "Metal" : "CPU"); - return false; - } - - if (!dump_named_tensors(*state, names, dump_dir, result.infos)) { - return false; - } - - state.reset(); - sam3_free_model(*model); - model.reset(); - return true; -} - -static ggml_backend_t create_backend(bool use_gpu, int n_threads) { -#ifdef GGML_USE_METAL - if (use_gpu) { - return ggml_backend_metal_init(); - } -#else - (void) use_gpu; -#endif - ggml_backend_t backend = ggml_backend_cpu_init(); - if (backend) { - ggml_backend_cpu_set_n_threads(backend, n_threads); - } - return backend; -} - -static std::vector normalize_shape_4d(const std::vector & shape) { - std::vector out = shape; - while (out.size() < 4) { - out.push_back(1); - } - return out; -} - -static ggml_tensor * new_tensor_4d_from_shape(ggml_context * ctx, - ggml_type type, - const std::vector & shape) { - const std::vector s = normalize_shape_4d(shape); - if (s.size() != 4) { - return nullptr; - } - return ggml_new_tensor_4d(ctx, type, s[0], s[1], s[2], s[3]); -} - -static isolated_run_result run_isolated_attention(ggml_backend_t backend, - const tensor_case & tc, - bool force_contiguous_v) { - isolated_run_result result; - - const size_t ctx_size = ggml_tensor_overhead() * 16 + ggml_graph_overhead(); - ggml_init_params params = { - /*.mem_size =*/ ctx_size, - /*.mem_buffer =*/ nullptr, - /*.no_alloc =*/ true, - }; - ggml_context * ctx = ggml_init(params); - if (!ctx) { - return result; - } - - ggml_tensor * q = new_tensor_4d_from_shape(ctx, GGML_TYPE_F32, tc.q_flash.shape); - ggml_tensor * k = new_tensor_4d_from_shape(ctx, GGML_TYPE_F32, tc.k_flash.shape); - ggml_tensor * v_src = nullptr; - ggml_tensor * v = nullptr; - - ggml_set_input(q); - ggml_set_input(k); - if (!q || !k) { - ggml_free(ctx); - return result; - } - - if (force_contiguous_v) { - v = new_tensor_4d_from_shape(ctx, GGML_TYPE_F32, tc.v_flash.shape); - ggml_set_input(v); - } else { - v_src = new_tensor_4d_from_shape(ctx, GGML_TYPE_F32, tc.v_base.shape); - ggml_set_input(v_src); - v = ggml_permute(ctx, v_src, 0, 2, 1, 3); - } - if (!v) { - ggml_free(ctx); - return result; - } - - const float scale = 1.0f / std::sqrt((float) tc.q_flash.shape[0]); - ggml_tensor * out = ggml_flash_attn_ext(ctx, q, k, v, nullptr, scale, 0.0f, 0.0f); - ggml_set_output(out); - - ggml_cgraph * graph = ggml_new_graph(ctx); - ggml_build_forward_expand(graph, out); - - ggml_gallocr_t galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!ggml_gallocr_reserve(galloc, graph) || !ggml_gallocr_alloc_graph(galloc, graph)) { - ggml_gallocr_free(galloc); - ggml_free(ctx); - return result; - } - - ggml_backend_tensor_set(q, tc.q_flash.data.data(), 0, tc.q_flash.data.size() * sizeof(float)); - ggml_backend_tensor_set(k, tc.k_flash.data.data(), 0, tc.k_flash.data.size() * sizeof(float)); - if (force_contiguous_v) { - ggml_backend_tensor_set(v, tc.v_flash.data.data(), 0, tc.v_flash.data.size() * sizeof(float)); - } else { - ggml_backend_tensor_set(v_src, tc.v_base.data.data(), 0, tc.v_base.data.size() * sizeof(float)); - } - - auto t0 = std::chrono::high_resolution_clock::now(); - const ggml_status status = ggml_backend_graph_compute(backend, graph); - auto t1 = std::chrono::high_resolution_clock::now(); - result.elapsed_ms = std::chrono::duration(t1 - t0).count(); - if (status == GGML_STATUS_SUCCESS) { - result.shape = { - (int) out->ne[0], - (int) out->ne[1], - (int) out->ne[2], - (int) out->ne[3], - }; - result.output.resize((size_t) ggml_nelements(out)); - ggml_backend_tensor_get(out, result.output.data(), 0, result.output.size() * sizeof(float)); - } else { - fprintf(stderr, "isolated attention failed: %s\n", ggml_status_to_string(status)); - } - - ggml_gallocr_free(galloc); - ggml_free(ctx); - return result; -} - -static ref_tensor_f32 tensor_to_ref(ggml_tensor * t) { - ref_tensor_f32 out; - out.shape = { - (int) t->ne[0], - (int) t->ne[1], - (int) t->ne[2], - (int) t->ne[3], - }; - while (!out.shape.empty() && out.shape.back() == 1) { - out.shape.pop_back(); - } - if (out.shape.empty()) { - out.shape.push_back(1); - } - - out.data.resize((size_t) ggml_nelements(t)); - ggml_backend_tensor_get(t, out.data.data(), 0, out.data.size() * sizeof(float)); - return out; -} - -static prep_run_result run_isolated_qkv_prep(ggml_backend_t backend, const tensor_case & tc) { - prep_run_result result; - - const int E = 1024; - const int NH = 16; - const int HD = 64; - const int W = 72; - const int H = 72; - const int N = W * H; - - const size_t ctx_size = ggml_tensor_overhead() * 32 + ggml_graph_overhead(); - ggml_init_params params = { - /*.mem_size =*/ ctx_size, - /*.mem_buffer =*/ nullptr, - /*.no_alloc =*/ true, - }; - ggml_context * ctx = ggml_init(params); - if (!ctx) { - return result; - } - - ggml_tensor * qkv = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 3 * E, W, H, 1); - ggml_set_input(qkv); - ggml_set_name(qkv, "qkv_proj"); - - ggml_tensor * cur = ggml_reshape_4d(ctx, qkv, E, 3, N, 1); - cur = ggml_cont(ctx, ggml_permute(ctx, cur, 0, 3, 1, 2)); - - ggml_tensor * q_split = ggml_view_3d(ctx, cur, E, N, 1, cur->nb[1], cur->nb[2], 0); - ggml_tensor * k_split = ggml_view_3d(ctx, cur, E, N, 1, cur->nb[1], cur->nb[2], 1 * cur->nb[3]); - ggml_tensor * v_split = ggml_view_3d(ctx, cur, E, N, 1, cur->nb[1], cur->nb[2], 2 * cur->nb[3]); - ggml_set_name(q_split, "q_split"); - ggml_set_name(k_split, "k_split"); - ggml_set_name(v_split, "v_split"); - - ggml_tensor * q_heads_base = ggml_reshape_4d(ctx, q_split, HD, NH, N, 1); - ggml_tensor * k_heads_base = ggml_reshape_4d(ctx, k_split, HD, NH, N, 1); - ggml_tensor * v_heads_base = ggml_reshape_4d(ctx, v_split, HD, NH, N, 1); - ggml_set_name(q_heads_base, "q_heads_base"); - ggml_set_name(k_heads_base, "k_heads_base"); - ggml_set_name(v_heads_base, "v_heads_base"); - - ggml_tensor * q_heads = ggml_cont(ctx, ggml_permute(ctx, q_heads_base, 0, 2, 1, 3)); - ggml_tensor * k_heads = ggml_cont(ctx, ggml_permute(ctx, k_heads_base, 0, 2, 1, 3)); - ggml_tensor * v_flash = ggml_permute(ctx, v_heads_base, 0, 2, 1, 3); - ggml_set_name(q_heads, "q_heads"); - ggml_set_name(k_heads, "k_heads"); - ggml_set_name(v_flash, "v_flash"); - - std::vector outs = { - q_split, - k_split, - v_split, - q_heads_base, - k_heads_base, - v_heads_base, - q_heads, - k_heads, - v_flash, - }; - ggml_cgraph * graph = ggml_new_graph(ctx); - for (ggml_tensor * out : outs) { - ggml_set_output(out); - ggml_build_forward_expand(graph, out); - } - - ggml_gallocr_t galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!ggml_gallocr_reserve(galloc, graph) || !ggml_gallocr_alloc_graph(galloc, graph)) { - ggml_gallocr_free(galloc); - ggml_free(ctx); - return result; - } - - ggml_backend_tensor_set(qkv, tc.qkv_proj.data.data(), 0, tc.qkv_proj.data.size() * sizeof(float)); - const ggml_status status = ggml_backend_graph_compute(backend, graph); - if (status == GGML_STATUS_SUCCESS) { - for (ggml_tensor * out : outs) { - result.tensors[out->name] = tensor_to_ref(out); - } - } else { - fprintf(stderr, "isolated qkv prep failed: %s\n", ggml_status_to_string(status)); - } - - ggml_gallocr_free(galloc); - ggml_free(ctx); - return result; -} - -static isolated_run_result run_isolated_rope(ggml_backend_t backend, - const ref_tensor_f32 & x_in, - const ref_tensor_f32 & freqs_in) { - isolated_run_result result; - - const std::vector x_shape = normalize_shape_4d(x_in.shape); - const std::vector freqs_shape = normalize_shape_4d(freqs_in.shape); - if (x_shape[0] != 64 || freqs_shape[0] != 2) { - return result; - } - - const int64_t head_dim = x_shape[0]; - const int64_t N = x_shape[1]; - const int64_t nheads_B = x_shape[2] * x_shape[3]; - const int64_t half = head_dim / 2; - - const size_t ctx_size = ggml_tensor_overhead() * 32 + ggml_graph_overhead(); - ggml_init_params params = { - /*.mem_size =*/ ctx_size, - /*.mem_buffer =*/ nullptr, - /*.no_alloc =*/ true, - }; - ggml_context * ctx = ggml_init(params); - if (!ctx) { - return result; - } - - ggml_tensor * x = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, x_shape[0], x_shape[1], x_shape[2] * x_shape[3]); - ggml_tensor * freqs = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, freqs_shape[0], freqs_shape[1], freqs_shape[2] * freqs_shape[3]); - ggml_set_input(x); - ggml_set_input(freqs); - - ggml_tensor * x_pairs = ggml_reshape_4d(ctx, x, 2, half, N, nheads_B); - ggml_tensor * fc = ggml_reshape_4d(ctx, freqs, 2, half, N, 1); - - ggml_tensor * cos_f = ggml_view_4d(ctx, fc, 1, half, N, 1, fc->nb[1], fc->nb[2], fc->nb[3], 0); - ggml_tensor * sin_f = ggml_view_4d(ctx, fc, 1, half, N, 1, fc->nb[1], fc->nb[2], fc->nb[3], fc->nb[0]); - ggml_tensor * x_re = ggml_view_4d(ctx, x_pairs, 1, half, N, nheads_B, x_pairs->nb[1], x_pairs->nb[2], x_pairs->nb[3], 0); - ggml_tensor * x_im = ggml_view_4d(ctx, x_pairs, 1, half, N, nheads_B, x_pairs->nb[1], x_pairs->nb[2], x_pairs->nb[3], x_pairs->nb[0]); - - ggml_tensor * out_re = ggml_sub(ctx, ggml_mul(ctx, x_re, cos_f), ggml_mul(ctx, x_im, sin_f)); - ggml_tensor * out_im = ggml_add(ctx, ggml_mul(ctx, x_re, sin_f), ggml_mul(ctx, x_im, cos_f)); - ggml_tensor * out = ggml_concat(ctx, out_re, out_im, 0); - out = ggml_reshape_3d(ctx, ggml_cont(ctx, out), head_dim, N, nheads_B); - ggml_set_output(out); - - ggml_cgraph * graph = ggml_new_graph(ctx); - ggml_build_forward_expand(graph, out); - - ggml_gallocr_t galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!ggml_gallocr_reserve(galloc, graph) || !ggml_gallocr_alloc_graph(galloc, graph)) { - ggml_gallocr_free(galloc); - ggml_free(ctx); - return result; - } - - ggml_backend_tensor_set(x, x_in.data.data(), 0, x_in.data.size() * sizeof(float)); - ggml_backend_tensor_set(freqs, freqs_in.data.data(), 0, freqs_in.data.size() * sizeof(float)); - const ggml_status status = ggml_backend_graph_compute(backend, graph); - if (status == GGML_STATUS_SUCCESS) { - result.shape = { (int) out->ne[0], (int) out->ne[1], (int) out->ne[2] }; - result.output.resize((size_t) ggml_nelements(out)); - ggml_backend_tensor_get(out, result.output.data(), 0, result.output.size() * sizeof(float)); - } else { - fprintf(stderr, "isolated rope failed: %s\n", ggml_status_to_string(status)); - } - - ggml_gallocr_free(galloc); - ggml_free(ctx); - return result; -} - -static bool load_tensor_case(const std::string & dir, tensor_case & tc) { - tc.qkv_proj = load_ref_f32(dir + "/dbg_block_15_qkv_proj"); - tc.q_heads = load_ref_f32(dir + "/dbg_block_15_q_heads"); - tc.k_heads = load_ref_f32(dir + "/dbg_block_15_k_heads"); - tc.q_rope = load_ref_f32(dir + "/dbg_block_15_q_rope"); - tc.k_rope = load_ref_f32(dir + "/dbg_block_15_k_rope"); - tc.q_flash = load_ref_f32(dir + "/dbg_block_15_q_flash"); - tc.k_flash = load_ref_f32(dir + "/dbg_block_15_k_flash"); - tc.v_base = load_ref_f32(dir + "/dbg_block_15_v_heads_base"); - tc.v_flash = load_ref_f32(dir + "/dbg_block_15_v_flash"); - tc.attn_out = load_ref_f32(dir + "/dbg_block_15_attn_out"); - - return !tc.qkv_proj.data.empty() && - !tc.q_heads.data.empty() && - !tc.k_heads.data.empty() && - !tc.q_rope.data.empty() && - !tc.k_rope.data.empty() && - !tc.q_flash.data.empty() && - !tc.k_flash.data.empty() && - !tc.v_base.data.empty() && - !tc.v_flash.data.empty() && - !tc.attn_out.data.empty(); -} - -static void print_tensor_info(const char * label, - const std::string & name, - const sam3_tensor_info & info) { - fprintf(stderr, - "%s %-24s shape=[%lld,%lld,%lld,%lld] strides=[%llu,%llu,%llu,%llu] type=%s contiguous=%d op=%d\n", - label, - name.c_str(), - (long long) info.ne[0], - (long long) info.ne[1], - (long long) info.ne[2], - (long long) info.ne[3], - (unsigned long long) info.nb[0], - (unsigned long long) info.nb[1], - (unsigned long long) info.nb[2], - (unsigned long long) info.nb[3], - ggml_type_name_public(info.type), - info.is_contiguous ? 1 : 0, - info.op); -} - -int main(int argc, char ** argv) { - if (argc < 2) { - usage(argv[0]); - return 1; - } - - const std::string model_path = argv[1]; - const std::string ref_dir = argc >= 3 ? argv[2] : "tests/ref_phase3"; - - auto preprocessed = load_ref_f32(ref_dir + "/preprocessed"); - if (preprocessed.data.empty() || preprocessed.shape.size() != 4) { - fprintf(stderr, "failed to load %s/preprocessed\n", ref_dir.c_str()); - return 1; - } - - const int img_size = preprocessed.shape[2]; - const std::vector names = { - "dbg_block_15_norm1", - "dbg_block_15_qkv_proj", - "dbg_block_15_q_split", - "dbg_block_15_k_split", - "dbg_block_15_v_split", - "dbg_block_15_q_heads_base", - "dbg_block_15_k_heads_base", - "dbg_block_15_v_heads_base", - "dbg_block_15_q_heads", - "dbg_block_15_k_heads", - "dbg_block_15_v_flash", - "dbg_block_15_q_rope", - "dbg_block_15_k_rope", - "dbg_block_15_q_flash", - "dbg_block_15_k_flash", - "dbg_block_15_attn_out", - }; - - const std::string cpu_dir = "/tmp/sam3_block15_flash_cpu"; - const std::string metal_dir = "/tmp/sam3_block15_flash_metal"; - ensure_dir(cpu_dir); - ensure_dir(metal_dir); - - backend_run_result cpu_run; - backend_run_result metal_run; - - fprintf(stderr, "\n=== Extract CPU Block-15 Attention Tensors ===\n"); - if (!run_encoder_backend(model_path, preprocessed.data, img_size, false, 8, names, cpu_dir, cpu_run)) { - return 1; - } - fprintf(stderr, "CPU extract wall time: %.1f ms\n", cpu_run.elapsed_ms); - - fprintf(stderr, "\n=== Extract Metal Block-15 Attention Tensors ===\n"); - if (!run_encoder_backend(model_path, preprocessed.data, img_size, true, 8, names, metal_dir, metal_run)) { - return 1; - } - fprintf(stderr, "Metal extract wall time: %.1f ms\n", metal_run.elapsed_ms); - - fprintf(stderr, "\n=== Exact Attention Case ===\n"); - fprintf(stderr, "block=15 variant=global_self_attention heads=16 head_dim=64 tokens=72*72=5184 batch=1 mask=none\n"); - print_tensor_info("CPU ", "dbg_block_15_q_flash", cpu_run.infos["dbg_block_15_q_flash"]); - print_tensor_info("CPU ", "dbg_block_15_k_flash", cpu_run.infos["dbg_block_15_k_flash"]); - print_tensor_info("CPU ", "dbg_block_15_v_heads_base", cpu_run.infos["dbg_block_15_v_heads_base"]); - print_tensor_info("CPU ", "dbg_block_15_v_flash", cpu_run.infos["dbg_block_15_v_flash"]); - print_tensor_info("Metal", "dbg_block_15_q_flash", metal_run.infos["dbg_block_15_q_flash"]); - print_tensor_info("Metal", "dbg_block_15_k_flash", metal_run.infos["dbg_block_15_k_flash"]); - print_tensor_info("Metal", "dbg_block_15_v_heads_base", metal_run.infos["dbg_block_15_v_heads_base"]); - print_tensor_info("Metal", "dbg_block_15_v_flash", metal_run.infos["dbg_block_15_v_flash"]); - - fprintf(stderr, "\n=== CPU vs Metal Attention Prep Checkpoints ===\n"); - fprintf(stderr, "%-24s %-18s %14s %14s %12s\n", - "checkpoint", "shape", "max_abs_diff", "mean_abs_diff", "n_bad"); - for (const auto & name : names) { - auto cpu = load_ref_f32(cpu_dir + "/" + name); - auto metal = load_ref_f32(metal_dir + "/" + name); - if (cpu.data.empty() || metal.data.empty() || cpu.shape != metal.shape) { - fprintf(stderr, "%-24s %-18s %14s %14s %12s\n", - name.c_str(), "-", "load-fail", "load-fail", "-"); - continue; - } - compare_result diff = compare_tensors(cpu.data.data(), metal.data.data(), cpu.numel(), 1e-4f); - fprintf(stderr, "%-24s %-18s %14.6f %14.6f %12d\n", - name.c_str(), format_shape(cpu.shape).c_str(), diff.max_diff, diff.mean_diff, diff.n_bad); - } - - tensor_case tc; - if (!load_tensor_case(cpu_dir, tc)) { - fprintf(stderr, "failed to load isolated attention tensors from %s\n", cpu_dir.c_str()); - return 1; - } - - { - sam3_params params; - params.model_path = model_path; - params.use_gpu = false; - params.n_threads = 8; - auto model = sam3_load_model(params); - if (!model) { - fprintf(stderr, "failed to reload model for freqs_cis dump\n"); - return 1; - } - const std::string freqs_path = "/tmp/sam3_block15_flash_cpu/vit.blocks.15.attn.freqs_cis"; - if (!sam3_dump_model_tensor(*model, "vit.blocks.15.attn.freqs_cis", freqs_path)) { - fprintf(stderr, "failed to dump vit.blocks.15.attn.freqs_cis\n"); - sam3_free_model(*model); - model.reset(); - return 1; - } - tc.freqs_cis = load_ref_f32(freqs_path); - sam3_free_model(*model); - model.reset(); - if (tc.freqs_cis.data.empty()) { - fprintf(stderr, "failed to load dumped freqs_cis\n"); - return 1; - } - } - - ggml_backend_t cpu_backend = create_backend(false, 8); - ggml_backend_t metal_backend = create_backend(true, 8); - if (!cpu_backend || !metal_backend) { - fprintf(stderr, "failed to initialize CPU or Metal backend\n"); - if (cpu_backend) ggml_backend_free(cpu_backend); - if (metal_backend) ggml_backend_free(metal_backend); - return 1; - } - - fprintf(stderr, "\n=== Isolated QKV Preparation From Exact qkv_proj ===\n"); - prep_run_result cpu_prep = run_isolated_qkv_prep(cpu_backend, tc); - prep_run_result metal_prep = run_isolated_qkv_prep(metal_backend, tc); - const char * prep_names[] = { - "q_split", - "k_split", - "v_split", - "q_heads_base", - "k_heads_base", - "v_heads_base", - "q_heads", - "k_heads", - "v_flash", - }; - fprintf(stderr, "%-16s %-18s %14s %14s %12s %14s\n", - "stage", "shape", "cpu_vs_dump", "metal_vs_cpu", "n_bad", "mean_diff"); - for (const char * stage : prep_names) { - const auto cpu_it = cpu_prep.tensors.find(stage); - const auto metal_it = metal_prep.tensors.find(stage); - if (cpu_it == cpu_prep.tensors.end() || metal_it == metal_prep.tensors.end()) { - fprintf(stderr, "%-16s %-18s %14s %14s %12s %14s\n", - stage, "-", "missing", "missing", "-", "-"); - continue; - } - - const std::string dump_name = std::string("dbg_block_15_") + stage; - auto cpu_dump = load_ref_f32(cpu_dir + "/" + dump_name); - compare_result cpu_vs_dump_stage = compare_tensors(cpu_it->second.data.data(), cpu_dump.data.data(), cpu_it->second.numel(), 1e-4f); - compare_result metal_vs_cpu_stage = compare_tensors(metal_it->second.data.data(), cpu_it->second.data.data(), cpu_it->second.numel(), 1e-4f); - fprintf(stderr, "%-16s %-18s %14.6f %14.6f %12d %14.6f\n", - stage, - format_shape(cpu_it->second.shape).c_str(), - cpu_vs_dump_stage.max_diff, - metal_vs_cpu_stage.max_diff, - metal_vs_cpu_stage.n_bad, - metal_vs_cpu_stage.mean_diff); - } - - fprintf(stderr, "\n=== Isolated RoPE On Exact q_heads / k_heads ===\n"); - isolated_run_result cpu_q_rope = run_isolated_rope(cpu_backend, tc.q_heads, tc.freqs_cis); - isolated_run_result metal_q_rope = run_isolated_rope(metal_backend, tc.q_heads, tc.freqs_cis); - isolated_run_result cpu_k_rope = run_isolated_rope(cpu_backend, tc.k_heads, tc.freqs_cis); - isolated_run_result metal_k_rope = run_isolated_rope(metal_backend, tc.k_heads, tc.freqs_cis); - compare_result q_rope_cpu_vs_dump = compare_tensors(cpu_q_rope.output.data(), tc.q_rope.data.data(), (int) tc.q_rope.data.size(), 1e-4f); - compare_result q_rope_metal_vs_cpu = compare_tensors(metal_q_rope.output.data(), cpu_q_rope.output.data(), (int) cpu_q_rope.output.size(), 1e-4f); - compare_result k_rope_cpu_vs_dump = compare_tensors(cpu_k_rope.output.data(), tc.k_rope.data.data(), (int) tc.k_rope.data.size(), 1e-4f); - compare_result k_rope_metal_vs_cpu = compare_tensors(metal_k_rope.output.data(), cpu_k_rope.output.data(), (int) cpu_k_rope.output.size(), 1e-4f); - fprintf(stderr, - "q_rope cpu_vs_dump(max=%.6f mean=%.6f) metal_vs_cpu(max=%.6f mean=%.6f n_bad=%d)\n", - q_rope_cpu_vs_dump.max_diff, - q_rope_cpu_vs_dump.mean_diff, - q_rope_metal_vs_cpu.max_diff, - q_rope_metal_vs_cpu.mean_diff, - q_rope_metal_vs_cpu.n_bad); - fprintf(stderr, - "k_rope cpu_vs_dump(max=%.6f mean=%.6f) metal_vs_cpu(max=%.6f mean=%.6f n_bad=%d)\n", - k_rope_cpu_vs_dump.max_diff, - k_rope_cpu_vs_dump.mean_diff, - k_rope_metal_vs_cpu.max_diff, - k_rope_metal_vs_cpu.mean_diff, - k_rope_metal_vs_cpu.n_bad); - - fprintf(stderr, "\n=== Isolated FLASH_ATTN_EXT (Exact V View) ===\n"); - isolated_run_result cpu_exact = run_isolated_attention(cpu_backend, tc, false); - isolated_run_result metal_exact = run_isolated_attention(metal_backend, tc, false); - if (cpu_exact.output.empty() || metal_exact.output.empty()) { - fprintf(stderr, "isolated exact-view attention did not produce outputs\n"); - ggml_backend_free(cpu_backend); - ggml_backend_free(metal_backend); - return 1; - } - compare_result cpu_vs_dump = compare_tensors(cpu_exact.output.data(), tc.attn_out.data.data(), (int) tc.attn_out.data.size(), 1e-4f); - compare_result metal_vs_cpu = compare_tensors(metal_exact.output.data(), cpu_exact.output.data(), (int) cpu_exact.output.size(), 1e-4f); - fprintf(stderr, - "exact_view shape=%s cpu_ms=%.3f metal_ms=%.3f cpu_vs_dump(max=%.6f mean=%.6f) metal_vs_cpu(max=%.6f mean=%.6f n_bad=%d)\n", - format_shape(cpu_exact.shape).c_str(), - cpu_exact.elapsed_ms, - metal_exact.elapsed_ms, - cpu_vs_dump.max_diff, - cpu_vs_dump.mean_diff, - metal_vs_cpu.max_diff, - metal_vs_cpu.mean_diff, - metal_vs_cpu.n_bad); - - fprintf(stderr, "\n=== Isolated FLASH_ATTN_EXT (Forced Contiguous V) ===\n"); - isolated_run_result cpu_cont = run_isolated_attention(cpu_backend, tc, true); - isolated_run_result metal_cont = run_isolated_attention(metal_backend, tc, true); - if (cpu_cont.output.empty() || metal_cont.output.empty()) { - fprintf(stderr, "isolated contiguous-V attention did not produce outputs\n"); - ggml_backend_free(cpu_backend); - ggml_backend_free(metal_backend); - return 1; - } - compare_result cont_cpu_vs_exact = compare_tensors(cpu_cont.output.data(), cpu_exact.output.data(), (int) cpu_exact.output.size(), 1e-4f); - compare_result cont_metal_vs_cpu = compare_tensors(metal_cont.output.data(), cpu_cont.output.data(), (int) cpu_cont.output.size(), 1e-4f); - fprintf(stderr, - "contiguous_v shape=%s cpu_ms=%.3f metal_ms=%.3f cpu_cont_vs_exact(max=%.6f mean=%.6f) metal_vs_cpu(max=%.6f mean=%.6f n_bad=%d)\n", - format_shape(cpu_cont.shape).c_str(), - cpu_cont.elapsed_ms, - metal_cont.elapsed_ms, - cont_cpu_vs_exact.max_diff, - cont_cpu_vs_exact.mean_diff, - cont_metal_vs_cpu.max_diff, - cont_metal_vs_cpu.mean_diff, - cont_metal_vs_cpu.n_bad); - - ggml_backend_free(cpu_backend); - ggml_backend_free(metal_backend); - return 0; -} diff --git a/tests/test_metal_overhead.cpp b/tests/test_metal_overhead.cpp deleted file mode 100644 index 8ea99c2..0000000 --- a/tests/test_metal_overhead.cpp +++ /dev/null @@ -1,377 +0,0 @@ -// Measure Metal dispatch overhead and its impact on SAM3 ViT performance. -// Proves WHY Metal is slower than CPU for the full 2524-node ViT graph. - -#include "ggml-alloc.h" -#include "ggml-backend.h" -#include "ggml-cpu.h" -#include "ggml.h" - -#ifdef GGML_USE_METAL -#include "ggml-metal.h" -#endif - -#include -#include -#include -#include - -static double run_graph(ggml_backend_t backend, ggml_context* ctx, ggml_cgraph* graph, - const std::vector& inputs, int warmup, int repeats) { - auto* galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!ggml_gallocr_reserve(galloc, graph) || !ggml_gallocr_alloc_graph(galloc, graph)) { - fprintf(stderr, "FAIL: alloc\n"); - ggml_gallocr_free(galloc); - return -1; - } - for (auto* t : inputs) { - int64_t nel = ggml_nelements(t); - if (t->type == GGML_TYPE_F32) { - std::vector d(nel, 0.1f); - ggml_backend_tensor_set(t, d.data(), 0, nel * sizeof(float)); - } else if (t->type == GGML_TYPE_F16) { - std::vector d(nel); - for (int64_t j = 0; j < nel; j++) d[j] = ggml_fp32_to_fp16(0.1f); - ggml_backend_tensor_set(t, d.data(), 0, nel * sizeof(uint16_t)); - } - } - if (ggml_backend_is_cpu(backend)) - ggml_backend_cpu_set_n_threads(backend, 8); - - for (int i = 0; i < warmup; i++) - ggml_backend_graph_compute(backend, graph); - auto t0 = std::chrono::high_resolution_clock::now(); - for (int i = 0; i < repeats; i++) - ggml_backend_graph_compute(backend, graph); - auto t1 = std::chrono::high_resolution_clock::now(); - ggml_gallocr_free(galloc); - return std::chrono::duration(t1 - t0).count() / repeats; -} - -// ═══════════════════════════════════════════════════════════════════════════════ -// TEST 1: Pure dispatch overhead — chain of N trivial add ops -// ═══════════════════════════════════════════════════════════════════════════════ -static void test_dispatch_overhead(ggml_backend_t cpu, ggml_backend_t metal) { - fprintf(stderr, "\n=== TEST 1: Dispatch overhead (chain of N add ops on [1024,72,72] tensors) ===\n\n"); - fprintf(stderr, " %6s %10s %10s %12s\n", "N_ops", "CPU (ms)", "Metal (ms)", "Overhead/op"); - - int counts[] = {1, 10, 50, 100, 250, 500, 1000, 2500}; - double metal_1op = 0; - - for (int N : counts) { - int max_tensors = N + 20; - int max_graph = N + 20; - - auto build_and_run = [&](ggml_backend_t backend) -> double { - size_t ctx_size = ggml_tensor_overhead() * max_tensors + ggml_graph_overhead_custom(max_graph, false); - ggml_init_params params = {ctx_size, nullptr, true}; - auto* ctx = ggml_init(params); - - auto* a = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1024, 72, 72); ggml_set_input(a); - auto* b = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1024, 72, 72); ggml_set_input(b); - auto* x = ggml_add(ctx, a, b); - for (int i = 1; i < N; i++) - x = ggml_add(ctx, x, b); - ggml_set_output(x); - - auto* graph = ggml_new_graph_custom(ctx, max_graph, false); - ggml_build_forward_expand(graph, x); - - std::vector inputs = {a, b}; - double ms = run_graph(backend, ctx, graph, inputs, 1, 3); - ggml_free(ctx); - return ms; - }; - - double cpu_ms = build_and_run(cpu); - double metal_ms = build_and_run(metal); - - if (N == 1) metal_1op = metal_ms; - double overhead_per_op = (N > 1) ? (metal_ms - metal_1op) / (N - 1) : 0; - - fprintf(stderr, " %6d %8.2f ms %8.2f ms %10.3f ms\n", - N, cpu_ms, metal_ms, overhead_per_op); - } -} - -// ═══════════════════════════════════════════════════════════════════════════════ -// TEST 2: Single ViT block as ONE graph vs sum of individual ops -// ═══════════════════════════════════════════════════════════════════════════════ -static void test_vit_block_graph(ggml_backend_t cpu, ggml_backend_t metal) { - fprintf(stderr, "\n=== TEST 2: One ViT block as single graph (windowed, all ops together) ===\n\n"); - - const int E = 1024, HD = 64, NH = 16, MLP = 4736; - const int W = 24, H = 24, NP = 9; // windowed - const ggml_type WT = GGML_TYPE_F16; - - auto run_block = [&](ggml_backend_t backend) -> std::pair> { - size_t ctx_size = ggml_tensor_overhead() * 128 + ggml_graph_overhead_custom(128, false); - ggml_init_params params = {ctx_size, nullptr, true}; - auto* ctx = ggml_init(params); - std::vector inputs; - - auto inp = [&](ggml_type type, int64_t d0, int64_t d1=1, int64_t d2=1, int64_t d3=1) { - auto* t = ggml_new_tensor_4d(ctx, type, d0, d1, d2, d3); - ggml_set_input(t); - inputs.push_back(t); - return t; - }; - - // Input + residual - auto* x = inp(GGML_TYPE_F32, E, W, H, NP); - - // --- Norm1 --- - auto* norm1_w = inp(GGML_TYPE_F32, E); - auto* norm1_b = inp(GGML_TYPE_F32, E); - auto* xn = ggml_add(ctx, ggml_mul(ctx, ggml_norm(ctx, x, 1e-6f), norm1_w), norm1_b); - - // --- QKV matmul --- - auto* qkv_w = inp(WT, E, 3*E); - auto* qkv_b = inp(GGML_TYPE_F32, 3*E); - auto* qkv = ggml_add(ctx, ggml_mul_mat(ctx, qkv_w, xn), qkv_b); - - // --- Reshape + split Q/K/V --- - int N = W * H; - qkv = ggml_reshape_4d(ctx, qkv, E, 3, N, NP); - qkv = ggml_cont(ctx, ggml_permute(ctx, qkv, 0, 3, 1, 2)); // expensive cont #1 - auto* Q = ggml_view_3d(ctx, qkv, E, N, NP, qkv->nb[1], qkv->nb[2], 0); - auto* K = ggml_view_3d(ctx, qkv, E, N, NP, qkv->nb[1], qkv->nb[2], qkv->nb[3]); - auto* V = ggml_view_3d(ctx, qkv, E, N, NP, qkv->nb[1], qkv->nb[2], 2*qkv->nb[3]); - - // --- Multi-head reshape + permute + cont (×3) --- - Q = ggml_reshape_4d(ctx, Q, HD, NH, N, NP); - Q = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 1, 3)); // expensive cont #2 - Q = ggml_reshape_4d(ctx, Q, HD, N, NH, NP); - - K = ggml_reshape_4d(ctx, K, HD, NH, N, NP); - K = ggml_cont(ctx, ggml_permute(ctx, K, 0, 2, 1, 3)); // expensive cont #3 - K = ggml_reshape_4d(ctx, K, HD, N, NH, NP); - - V = ggml_reshape_4d(ctx, V, HD, NH, N, NP); - V = ggml_permute(ctx, V, 0, 2, 1, 3); // non-contiguous view — flash_attn uses strides - - // --- Flash attention --- - auto* attn = ggml_flash_attn_ext(ctx, Q, K, V, nullptr, 1.0f/8.0f, 0, 0); - auto* attn_r = ggml_reshape_4d(ctx, attn, E, W, H, NP); - - // --- Output proj --- - auto* proj_w = inp(WT, E, E); - auto* proj_b = inp(GGML_TYPE_F32, E); - auto* proj = ggml_add(ctx, ggml_mul_mat(ctx, proj_w, attn_r), proj_b); - - // --- Residual --- - auto* res1 = ggml_add(ctx, x, proj); - - // --- Norm2 --- - auto* norm2_w = inp(GGML_TYPE_F32, E); - auto* norm2_b = inp(GGML_TYPE_F32, E); - auto* xn2 = ggml_add(ctx, ggml_mul(ctx, ggml_norm(ctx, res1, 1e-6f), norm2_w), norm2_b); - - // --- MLP --- - auto* fc1_w = inp(WT, E, MLP); - auto* fc1_b = inp(GGML_TYPE_F32, MLP); - auto* fc2_w = inp(WT, MLP, E); - auto* fc2_b = inp(GGML_TYPE_F32, E); - - auto* h = ggml_add(ctx, ggml_mul_mat(ctx, fc1_w, xn2), fc1_b); - h = ggml_gelu_erf(ctx, h); - h = ggml_add(ctx, ggml_mul_mat(ctx, fc2_w, h), fc2_b); - - // --- Residual --- - auto* out = ggml_add(ctx, res1, h); - ggml_set_output(out); - - auto* graph = ggml_new_graph_custom(ctx, 128, false); - ggml_build_forward_expand(graph, out); - - int n_nodes = ggml_graph_n_nodes(graph); - - // Allocate and fill inputs - auto* galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - ggml_gallocr_reserve(galloc, graph); - ggml_gallocr_alloc_graph(galloc, graph); - - // Fill inputs with deterministic data - for (auto* t : inputs) { - int64_t nel = ggml_nelements(t); - if (t->type == GGML_TYPE_F32) { - std::vector d(nel); - for (int64_t j = 0; j < nel; j++) d[j] = 0.01f * ((j % 97) - 48); - ggml_backend_tensor_set(t, d.data(), 0, nel * sizeof(float)); - } else { - std::vector d(nel); - for (int64_t j = 0; j < nel; j++) d[j] = ggml_fp32_to_fp16(0.01f * ((j % 97) - 48)); - ggml_backend_tensor_set(t, d.data(), 0, nel * sizeof(uint16_t)); - } - } - - if (ggml_backend_is_cpu(backend)) - ggml_backend_cpu_set_n_threads(backend, 8); - - // Warmup + timed run - ggml_backend_graph_compute(backend, graph); - auto t0 = std::chrono::high_resolution_clock::now(); - for (int i = 0; i < 5; i++) - ggml_backend_graph_compute(backend, graph); - auto t1 = std::chrono::high_resolution_clock::now(); - double ms = std::chrono::duration(t1 - t0).count() / 5; - - // Read output - int64_t nel = ggml_nelements(out); - std::vector output(nel); - ggml_backend_tensor_get(out, output.data(), 0, nel * sizeof(float)); - - ggml_gallocr_free(galloc); - ggml_free(ctx); - fprintf(stderr, " %-20s %7.2f ms (%d nodes)\n", - ggml_backend_is_cpu(backend) ? "CPU" : "Metal", ms, n_nodes); - return std::make_pair(ms, output); - }; - - auto cpu_result = run_block(cpu); - auto metal_result = run_block(metal); - double cpu_block = cpu_result.first; - double metal_block = metal_result.first; - const auto& cpu_out = cpu_result.second; - const auto& metal_out = metal_result.second; - - // Compare outputs - float max_diff = 0; - int max_idx = 0; - for (size_t i = 0; i < cpu_out.size(); i++) { - float d = std::fabs(cpu_out[i] - metal_out[i]); - if (d > max_diff) { max_diff = d; max_idx = (int)i; } - } - fprintf(stderr, " Output comparison: max_diff=%.6f at [%d] (cpu=%.6f metal=%.6f) — %s\n", - max_diff, max_idx, cpu_out[max_idx], metal_out[max_idx], - max_diff < 0.01f ? "PASS" : "FAIL"); - - fprintf(stderr, " Speedup: %.2fx\n", cpu_block / metal_block); - fprintf(stderr, "\n Extrapolated to 32 layers:\n"); - fprintf(stderr, " CPU: %.0f ms\n", cpu_block * 32); - fprintf(stderr, " Metal: %.0f ms\n", metal_block * 32); - fprintf(stderr, " Speedup: %.2fx\n", (cpu_block * 32) / (metal_block * 32)); -} - -// ═══════════════════════════════════════════════════════════════════════════════ -// TEST 3: Scale from 1 to 32 ViT blocks in ONE graph -// ═══════════════════════════════════════════════════════════════════════════════ -static void test_scaling(ggml_backend_t cpu, ggml_backend_t metal) { - fprintf(stderr, "\n=== TEST 3: Scaling — N chained matmul+add blocks in one graph ===\n"); - fprintf(stderr, " (Each block: F16 matmul [1024,1024]x[1024,24,24,9] + add + norm + gelu)\n\n"); - fprintf(stderr, " %6s %6s %10s %10s %8s\n", "Blocks", "Nodes", "CPU (ms)", "Metal (ms)", "Speedup"); - - const int E = 1024; - const ggml_type WT = GGML_TYPE_F16; - int block_counts[] = {1, 4, 8, 16, 32}; - - for (int nblocks : block_counts) { - int max_nodes = nblocks * 20 + 20; - size_t ctx_size = ggml_tensor_overhead() * (max_nodes + 50) + - ggml_graph_overhead_custom(max_nodes + 50, false); - ggml_init_params params = {ctx_size, nullptr, true}; - - auto run_on = [&](ggml_backend_t backend) -> std::pair { - auto* ctx = ggml_init(params); - std::vector inputs; - - auto inp = [&](ggml_type type, int64_t d0, int64_t d1=1, int64_t d2=1, int64_t d3=1) { - auto* t = ggml_new_tensor_4d(ctx, type, d0, d1, d2, d3); - ggml_set_input(t); - inputs.push_back(t); - return t; - }; - - auto* x = inp(GGML_TYPE_F32, E, 24, 24, 9); - - for (int i = 0; i < nblocks; i++) { - auto* w = inp(WT, E, E); - auto* b = inp(GGML_TYPE_F32, E); - auto* nw = inp(GGML_TYPE_F32, E); - auto* nb_ = inp(GGML_TYPE_F32, E); - x = ggml_add(ctx, ggml_mul_mat(ctx, w, x), b); - x = ggml_add(ctx, ggml_mul(ctx, ggml_norm(ctx, x, 1e-6f), nw), nb_); - x = ggml_gelu_erf(ctx, x); - } - ggml_set_output(x); - - auto* graph = ggml_new_graph_custom(ctx, max_nodes + 50, false); - ggml_build_forward_expand(graph, x); - int n_nodes = ggml_graph_n_nodes(graph); - - double ms = run_graph(backend, ctx, graph, inputs, 2, 5); - ggml_free(ctx); - return {ms, n_nodes}; - }; - - double cpu_ms; int cpu_nodes; - { auto r = run_on(cpu); cpu_ms = r.first; cpu_nodes = r.second; } - double metal_ms; int metal_nodes; - { auto r = run_on(metal); metal_ms = r.first; metal_nodes = r.second; } - fprintf(stderr, " %6d %6d %8.2f ms %8.2f ms %7.2fx\n", - nblocks, cpu_nodes, cpu_ms, metal_ms, cpu_ms / metal_ms); - } -} - -// ═══════════════════════════════════════════════════════════════════════════════ -// TEST 4: cont(permute) cost in isolation vs inside a larger graph -// ═══════════════════════════════════════════════════════════════════════════════ -static void test_cont_in_graph(ggml_backend_t cpu, ggml_backend_t metal) { - fprintf(stderr, "\n=== TEST 4: N cont(permute) ops in ONE graph ===\n"); - fprintf(stderr, " (Measures whether cont overhead compounds in large graphs)\n\n"); - fprintf(stderr, " %6s %10s %10s %10s\n", "N_cont", "CPU (ms)", "Metal (ms)", "Metal/op"); - - int counts[] = {1, 4, 16, 32, 64, 96}; - - for (int N : counts) { - int max_nodes = N * 3 + 10; - size_t ctx_size = ggml_tensor_overhead() * (max_nodes + 20) + - ggml_graph_overhead_custom(max_nodes + 20, false); - ggml_init_params params = {ctx_size, nullptr, true}; - - auto run_on = [&](ggml_backend_t backend) -> double { - auto* ctx = ggml_init(params); - std::vector inputs; - - // Chain: permute → cont → permute → cont → ... - auto* x = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1024, 576, 144); - ggml_set_input(x); - inputs.push_back(x); - - ggml_tensor* last = x; - for (int i = 0; i < N; i++) { - last = ggml_cont(ctx, ggml_permute(ctx, last, 0, 2, 1, 3)); - } - ggml_set_output(last); - - auto* graph = ggml_new_graph_custom(ctx, max_nodes + 20, false); - ggml_build_forward_expand(graph, last); - double ms = run_graph(backend, ctx, graph, inputs, 2, 3); - ggml_free(ctx); - return ms; - }; - - double cpu_ms = run_on(cpu); - double metal_ms = run_on(metal); - fprintf(stderr, " %6d %8.2f ms %8.2f ms %8.2f ms\n", - N, cpu_ms, metal_ms, metal_ms / N); - } -} - -int main() { -#ifndef GGML_USE_METAL - fprintf(stderr, "Metal not available\n"); return 1; -#else - auto* cpu = ggml_backend_cpu_init(); - auto* metal = ggml_backend_metal_init(); - if (!metal) { fprintf(stderr, "No Metal\n"); return 1; } - - test_dispatch_overhead(cpu, metal); - test_vit_block_graph(cpu, metal); - test_scaling(cpu, metal); - test_cont_in_graph(cpu, metal); - - ggml_backend_free(cpu); - ggml_backend_free(metal); - return 0; -#endif -} diff --git a/tests/test_metal_phase6.cpp b/tests/test_metal_phase6.cpp deleted file mode 100644 index 0ea1c37..0000000 --- a/tests/test_metal_phase6.cpp +++ /dev/null @@ -1,195 +0,0 @@ -#include "sam3.h" - -#include "test_utils.h" - -#include -#include -#include -#include - -struct phase6_case { - std::string id; - sam3_pvs_params params; -}; - -struct tensor_case { - const char * name; - const char * label; - float tol; -}; - -static std::vector parse_points(const std::string & field) { - std::vector pts; - if (field.empty()) { - return pts; - } - - size_t start = 0; - while (start < field.size()) { - size_t end = field.find('|', start); - if (end == std::string::npos) { - end = field.size(); - } - - const std::string part = field.substr(start, end - start); - const size_t colon = part.find(':'); - if (colon != std::string::npos) { - pts.push_back({std::stof(part.substr(0, colon)), std::stof(part.substr(colon + 1))}); - } - - start = end + 1; - } - - return pts; -} - -static bool parse_box(const std::string & field, sam3_box & box) { - if (field.empty()) { - return false; - } - - float vals[4]; - size_t start = 0; - for (int i = 0; i < 4; ++i) { - size_t end = field.find(':', start); - if (end == std::string::npos) { - end = field.size(); - } - vals[i] = std::stof(field.substr(start, end - start)); - start = end + 1; - } - - box = {vals[0], vals[1], vals[2], vals[3]}; - return true; -} - -static std::vector load_cases(const std::string & path) { - std::vector cases; - std::ifstream f(path); - std::string line; - while (std::getline(f, line)) { - if (line.empty()) { - continue; - } - - std::vector fields; - size_t start = 0; - while (start <= line.size()) { - size_t end = line.find('\t', start); - if (end == std::string::npos) { - end = line.size(); - } - fields.push_back(line.substr(start, end - start)); - start = end + 1; - if (end == line.size()) { - break; - } - } - while (fields.size() < 5) { - fields.emplace_back(); - } - - phase6_case tc; - tc.id = fields[0]; - tc.params.multimask = fields[1] == "1"; - tc.params.pos_points = parse_points(fields[2]); - tc.params.neg_points = parse_points(fields[3]); - tc.params.use_box = parse_box(fields[4], tc.params.box); - cases.push_back(std::move(tc)); - } - - return cases; -} - -static const std::vector & tensor_cases() { - static const std::vector cases = { - {"sam_dec_final_queries", "Decoder final queries", 6e-3f}, - {"sam_dec_masks", "Decoder mask logits", 2e-2f}, - {"sam_dec_iou", "Decoder IoU predictions", 2e-3f}, - {"sam_dec_obj_score", "Decoder object score logit", 3e-3f}, - {"sam_dec_sam_token", "Decoder SAM token", 6e-3f}, - }; - return cases; -} - -int main(int argc, char ** argv) { - if (argc < 5) { - fprintf(stderr, "Usage: %s \n", argv[0]); - return 1; - } - - const std::string ref_dir = argv[1]; - const std::string prephase_ref_dir = argv[2]; - const std::string model_path = argv[3]; - const std::string cases_path = argv[4]; - const std::string cpp_root = ref_dir + "/cpp_out_phase6_metal"; - - auto cases = load_cases(cases_path); - if (cases.empty()) { - fprintf(stderr, "No phase 6 cases loaded from %s\n", cases_path.c_str()); - return 1; - } - - if (!ensure_dir(cpp_root)) { - fprintf(stderr, "Failed to create %s\n", cpp_root.c_str()); - return 1; - } - - sam3_params params; - params.model_path = model_path; - params.n_threads = 1; - params.use_gpu = true; - - auto model = sam3_load_model(params); - if (!model) { - fprintf(stderr, "Failed to load model from %s\n", model_path.c_str()); - return 1; - } - - bool overall_ok = true; - for (const auto & tc : cases) { - fprintf(stderr, "\n=== Metal Phase 6 Case %s ===\n", tc.id.c_str()); - - const std::string case_ref_dir = ref_dir + "/" + tc.id; - const std::string case_cpp_dir = cpp_root + "/" + tc.id; - if (!ensure_dir(case_cpp_dir)) { - fprintf(stderr, "Failed to create %s\n", case_cpp_dir.c_str()); - overall_ok = false; - continue; - } - - if (!sam3_test_dump_phase6_from_ref_inputs(*model, prephase_ref_dir, tc.params, - case_cpp_dir, params.n_threads)) { - fprintf(stderr, " [FAIL] could not run Metal phase 6 dump for %s\n", tc.id.c_str()); - overall_ok = false; - continue; - } - - for (const auto & tensor : tensor_cases()) { - auto ref = load_ref_f32(case_ref_dir + "/" + tensor.name); - auto got = load_ref_f32(case_cpp_dir + "/" + tensor.name); - if (ref.data.empty() || got.data.empty() || ref.numel() != got.numel()) { - fprintf(stderr, " [FAIL] %s missing or shape mismatch\n", tensor.name); - overall_ok = false; - continue; - } - - auto cmp = compare_tensors(got.data.data(), ref.data.data(), ref.numel(), tensor.tol); - const bool ok = cmp.max_diff <= tensor.tol; - fprintf(stderr, - " %-22s max=%.6f mean=%.6f cos=%.8f tol=%.6f %s\n", - tensor.name, - cmp.max_diff, - cmp.mean_diff, - cmp.cosine_sim, - tensor.tol, - ok ? "PASS" : "FAIL"); - if (!ok) { - overall_ok = false; - } - } - } - - sam3_free_model(*model); - return overall_ok ? 0 : 1; -} diff --git a/tests/test_metal_profile.cpp b/tests/test_metal_profile.cpp deleted file mode 100644 index 730f403..0000000 --- a/tests/test_metal_profile.cpp +++ /dev/null @@ -1,215 +0,0 @@ -// Profile individual ggml ops at SAM3 ViT shapes on CPU vs Metal. -// Each op gets its own tiny graph — runs in seconds, not minutes. - -#include "ggml-alloc.h" -#include "ggml-backend.h" -#include "ggml-cpu.h" -#include "ggml.h" - -#ifdef GGML_USE_METAL -#include "ggml-metal.h" -#endif - -#include -#include -#include -#include -#include -#include - -using build_fn = std::function>(ggml_context*)>; - -static double run_bench(ggml_backend_t backend, build_fn build, - int warmup = 2, int repeats = 5) { - size_t ctx_size = ggml_tensor_overhead() * 64 + ggml_graph_overhead(); - struct ggml_init_params params = {ctx_size, nullptr, true}; - auto* ctx = ggml_init(params); - - auto result = build(ctx); - auto* out = result.first; - auto& inputs = result.second; - ggml_set_output(out); - - auto* graph = ggml_new_graph(ctx); - ggml_build_forward_expand(graph, out); - - auto* galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!ggml_gallocr_reserve(galloc, graph) || !ggml_gallocr_alloc_graph(galloc, graph)) { - fprintf(stderr, "FAIL: alloc\n"); - ggml_gallocr_free(galloc); ggml_free(ctx); - return -1; - } - - for (auto* t : inputs) { - int64_t nel = ggml_nelements(t); - if (t->type == GGML_TYPE_F32) { - std::vector d(nel, 0.1f); - ggml_backend_tensor_set(t, d.data(), 0, nel * sizeof(float)); - } else if (t->type == GGML_TYPE_F16) { - std::vector d(nel); - for (int64_t j = 0; j < nel; j++) d[j] = ggml_fp32_to_fp16(0.1f); - ggml_backend_tensor_set(t, d.data(), 0, nel * sizeof(uint16_t)); - } - } - - if (ggml_backend_is_cpu(backend)) - ggml_backend_cpu_set_n_threads(backend, 8); - - for (int i = 0; i < warmup; i++) - ggml_backend_graph_compute(backend, graph); - - auto t0 = std::chrono::high_resolution_clock::now(); - for (int i = 0; i < repeats; i++) - ggml_backend_graph_compute(backend, graph); - auto t1 = std::chrono::high_resolution_clock::now(); - - double ms = std::chrono::duration(t1 - t0).count() / repeats; - ggml_gallocr_free(galloc); - ggml_free(ctx); - return ms; -} - -struct op_bench { const char* name; int count; double cpu_ms; double metal_ms; }; - -// Helper macros for input tensor creation -#define INP_F32(...) [&]{ auto* t = ggml_new_tensor(__VA_ARGS__); ggml_set_name(t,"i"); ggml_set_input(t); return t; }() -#define INP(ctx, type, ...) [&]{ int64_t ne[] = {__VA_ARGS__}; auto* t = ggml_new_tensor(ctx, type, sizeof(ne)/sizeof(ne[0]), ne); ggml_set_input(t); return t; }() - -int main() { -#ifndef GGML_USE_METAL - fprintf(stderr, "Metal not available\n"); return 1; -#else - auto* cpu = ggml_backend_cpu_init(); - auto* metal = ggml_backend_metal_init(); - if (!metal) { fprintf(stderr, "No Metal\n"); return 1; } - - const int E = 1024, MLP = 4736, HD = 64, NH = 16; - const int W_win = 24, H_win = 24, NP = 9; - const int W_glb = 72, H_glb = 72; - const int N_win = W_win * H_win; - const int N_glb = W_glb * H_glb; - const ggml_type WT = GGML_TYPE_F16; - - std::vector results; - - auto bench = [&](const char* name, int count, build_fn build) { - double cpu_ms = run_bench(cpu, build); - double metal_ms = run_bench(metal, build); - results.push_back({name, count, cpu_ms, metal_ms}); - fprintf(stderr, " %-45s CPU:%7.2f ms Metal:%7.2f ms (x%d → CPU %5.0f / Metal %5.0f)\n", - name, cpu_ms, metal_ms, count, cpu_ms*count, metal_ms*count); - }; - - fprintf(stderr, "\n=== SAM3 ViT Op Profiling (F16 weights, F32 activations, CPU 8 threads) ===\n\n"); - - // ── Matmuls (windowed: 28 layers × 4 = 112) ──────────────────────── - bench("matmul QKV win [1024,3072]x[1024,24,24,9]", 28, [&](ggml_context* c) { - auto* w = ggml_new_tensor_2d(c, WT, E, 3*E); ggml_set_input(w); - auto* x = ggml_new_tensor_4d(c, GGML_TYPE_F32, E, W_win, H_win, NP); ggml_set_input(x); - return std::make_pair(ggml_mul_mat(c, w, x), std::vector{w, x}); - }); - bench("matmul proj win [1024,1024]x[1024,24,24,9]", 28, [&](ggml_context* c) { - auto* w = ggml_new_tensor_2d(c, WT, E, E); ggml_set_input(w); - auto* x = ggml_new_tensor_4d(c, GGML_TYPE_F32, E, W_win, H_win, NP); ggml_set_input(x); - return std::make_pair(ggml_mul_mat(c, w, x), std::vector{w, x}); - }); - bench("matmul fc1 win [1024,4736]x[1024,24,24,9]", 28, [&](ggml_context* c) { - auto* w = ggml_new_tensor_2d(c, WT, E, MLP); ggml_set_input(w); - auto* x = ggml_new_tensor_4d(c, GGML_TYPE_F32, E, W_win, H_win, NP); ggml_set_input(x); - return std::make_pair(ggml_mul_mat(c, w, x), std::vector{w, x}); - }); - bench("matmul fc2 win [4736,1024]x[4736,24,24,9]", 28, [&](ggml_context* c) { - auto* w = ggml_new_tensor_2d(c, WT, MLP, E); ggml_set_input(w); - auto* x = ggml_new_tensor_4d(c, GGML_TYPE_F32, MLP, W_win, H_win, NP); ggml_set_input(x); - return std::make_pair(ggml_mul_mat(c, w, x), std::vector{w, x}); - }); - - // ── Matmuls (global: 4 layers × 4 = 16) ──────────────────────────── - bench("matmul QKV glb [1024,3072]x[1024,72,72,1]", 4, [&](ggml_context* c) { - auto* w = ggml_new_tensor_2d(c, WT, E, 3*E); ggml_set_input(w); - auto* x = ggml_new_tensor_4d(c, GGML_TYPE_F32, E, W_glb, H_glb, 1); ggml_set_input(x); - return std::make_pair(ggml_mul_mat(c, w, x), std::vector{w, x}); - }); - bench("matmul fc1 glb [1024,4736]x[1024,72,72,1]", 4, [&](ggml_context* c) { - auto* w = ggml_new_tensor_2d(c, WT, E, MLP); ggml_set_input(w); - auto* x = ggml_new_tensor_4d(c, GGML_TYPE_F32, E, W_glb, H_glb, 1); ggml_set_input(x); - return std::make_pair(ggml_mul_mat(c, w, x), std::vector{w, x}); - }); - - // ── Flash attention ───────────────────────────────────────────────── - bench("flash_attn win HD=64 N=576 NH=16 B=9", 28, [&](ggml_context* c) { - auto* Q = ggml_new_tensor_4d(c, GGML_TYPE_F32, HD, N_win, NH, NP); ggml_set_input(Q); - auto* K = ggml_new_tensor_4d(c, GGML_TYPE_F32, HD, N_win, NH, NP); ggml_set_input(K); - auto* V = ggml_new_tensor_4d(c, GGML_TYPE_F32, HD, N_win, NH, NP); ggml_set_input(V); - return std::make_pair(ggml_flash_attn_ext(c, Q, K, V, nullptr, 1.0f/8.0f, 0, 0), - std::vector{Q, K, V}); - }); - bench("flash_attn glb HD=64 N=5184 NH=16 B=1", 4, [&](ggml_context* c) { - auto* Q = ggml_new_tensor_4d(c, GGML_TYPE_F32, HD, N_glb, NH, 1); ggml_set_input(Q); - auto* K = ggml_new_tensor_4d(c, GGML_TYPE_F32, HD, N_glb, NH, 1); ggml_set_input(K); - auto* V = ggml_new_tensor_4d(c, GGML_TYPE_F32, HD, N_glb, NH, 1); ggml_set_input(V); - return std::make_pair(ggml_flash_attn_ext(c, Q, K, V, nullptr, 1.0f/8.0f, 0, 0), - std::vector{Q, K, V}); - }); - - // ── Layer norm (norm + mul + add) ─────────────────────────────────── - bench("layer_norm [1024, 72, 72]", 64, [&](ggml_context* c) { - auto* x = ggml_new_tensor_3d(c, GGML_TYPE_F32, E, W_glb, H_glb); ggml_set_input(x); - auto* w = ggml_new_tensor_1d(c, GGML_TYPE_F32, E); ggml_set_input(w); - auto* b = ggml_new_tensor_1d(c, GGML_TYPE_F32, E); ggml_set_input(b); - auto* n = ggml_add(c, ggml_mul(c, ggml_norm(c, x, 1e-6f), w), b); - return std::make_pair(n, std::vector{x, w, b}); - }); - - // ── GELU ──────────────────────────────────────────────────────────── - bench("gelu_erf [4736, 24, 24, 9]", 28, [&](ggml_context* c) { - auto* x = ggml_new_tensor_4d(c, GGML_TYPE_F32, MLP, W_win, H_win, NP); ggml_set_input(x); - return std::make_pair(ggml_gelu_erf(c, x), std::vector{x}); - }); - - // ── Add ───────────────────────────────────────────────────────────── - bench("add [1024, 72, 72]", 128, [&](ggml_context* c) { - auto* a = ggml_new_tensor_3d(c, GGML_TYPE_F32, E, W_glb, H_glb); ggml_set_input(a); - auto* b = ggml_new_tensor_3d(c, GGML_TYPE_F32, E, W_glb, H_glb); ggml_set_input(b); - return std::make_pair(ggml_add(c, a, b), std::vector{a, b}); - }); - - // ── Win part / unpart ─────────────────────────────────────────────── - bench("win_part [1024,72,72] -> [1024,24,24,9]", 28, [&](ggml_context* c) { - auto* x = ggml_new_tensor_3d(c, GGML_TYPE_F32, E, W_glb, H_glb); ggml_set_input(x); - return std::make_pair(ggml_win_part(c, x, 24), std::vector{x}); - }); - bench("win_unpart [1024,24,24,9] -> [1024,72,72]", 28, [&](ggml_context* c) { - auto* x = ggml_new_tensor_4d(c, GGML_TYPE_F32, E, W_win, H_win, NP); ggml_set_input(x); - return std::make_pair(ggml_win_unpart(c, x, W_glb, H_glb, 24), std::vector{x}); - }); - - // ── Cont (permute + make contiguous — happens ~96x in ViT) ────────── - bench("cont(permute) [1024, 576, 144]", 96, [&](ggml_context* c) { - auto* x = ggml_new_tensor_3d(c, GGML_TYPE_F32, E, N_win, NH*NP); ggml_set_input(x); - return std::make_pair(ggml_cont(c, ggml_permute(c, x, 0, 2, 1, 3)), - std::vector{x}); - }); - - // ── Summary ───────────────────────────────────────────────────────── - fprintf(stderr, "\n=== Estimated Total ViT Cost ===\n\n"); - fprintf(stderr, " %-45s %10s %10s %8s\n", "Op", "CPU total", "Metal total", "Speedup"); - fprintf(stderr, " %-45s %10s %10s %8s\n", - "---------------------------------------------", "---------", "-----------", "-------"); - - double total_cpu = 0, total_metal = 0; - for (auto& r : results) { - double ct = r.cpu_ms * r.count; - double mt = r.metal_ms * r.count; - total_cpu += ct; - total_metal += mt; - fprintf(stderr, " %-45s %7.0f ms %9.0f ms %7.2fx\n", r.name, ct, mt, ct/mt); - } - fprintf(stderr, " %-45s %7.0f ms %9.0f ms %7.2fx\n", - "TOTAL (estimated ViT)", total_cpu, total_metal, total_cpu/total_metal); - - ggml_backend_free(cpu); - ggml_backend_free(metal); - return 0; -#endif -} diff --git a/tests/test_metal_python_ref_compare.cpp b/tests/test_metal_python_ref_compare.cpp deleted file mode 100644 index fa5cf29..0000000 --- a/tests/test_metal_python_ref_compare.cpp +++ /dev/null @@ -1,503 +0,0 @@ -#include "sam3.h" - -#include "test_utils.h" - -#include -#include -#include -#include -#include -#include -#include - -struct metric_row { - std::string tensor; - std::string shape; - float max_abs = 0.0f; - float mean_abs = 0.0f; - float mean_rel = 0.0f; - float cosine = 0.0f; - bool measured = false; - std::string note; -}; - -struct backend_run { - std::string backend; - double encode_ms = 0.0; - double final_ms = 0.0; - std::vector encoder_rows; - std::vector final_rows; -}; - -static std::string shape_to_string(const std::vector & shape) { - std::string s; - for (size_t i = 0; i < shape.size(); ++i) { - if (i > 0) s += "x"; - s += std::to_string(shape[i]); - } - return s; -} - -static metric_row compute_metrics(const std::string & tensor, - const float * got, - const float * ref, - int n, - const std::vector & ref_shape) { - metric_row row; - row.tensor = tensor; - row.shape = shape_to_string(ref_shape); - if (n <= 0) { - row.note = "empty"; - return row; - } - - double sum_abs = 0.0; - double sum_rel = 0.0; - double dot_ab = 0.0; - double dot_aa = 0.0; - double dot_bb = 0.0; - float max_abs = 0.0f; - - for (int i = 0; i < n; ++i) { - const float a = got[i]; - const float b = ref[i]; - const float d = fabsf(a - b); - sum_abs += d; - sum_rel += d / (fabsf(b) + 1e-8f); - if (d > max_abs) { - max_abs = d; - } - dot_ab += (double) a * (double) b; - dot_aa += (double) a * (double) a; - dot_bb += (double) b * (double) b; - } - - row.max_abs = max_abs; - row.mean_abs = (float) (sum_abs / n); - row.mean_rel = (float) (sum_rel / n); - const double denom = sqrt(dot_aa) * sqrt(dot_bb); - row.cosine = denom > 0.0 ? (float) (dot_ab / denom) : 0.0f; - row.measured = true; - return row; -} - -static std::vector ggml_dump_to_nhwc_flat(const ref_tensor_f32 & cpp) { - std::vector out; - if (cpp.shape.size() < 3) { - return out; - } - - const int E = cpp.shape[0]; - const int W = cpp.shape[1]; - const int H = cpp.shape[2]; - out.resize((size_t) E * W * H); - for (int h = 0; h < H; ++h) { - for (int w = 0; w < W; ++w) { - for (int e = 0; e < E; ++e) { - out[(size_t) h * W * E + (size_t) w * E + e] = - cpp.data[(size_t) e + (size_t) w * E + (size_t) h * E * W]; - } - } - } - return out; -} - -static std::vector ggml_dump_to_nchw_flat(const ref_tensor_f32 & cpp) { - std::vector out; - if (cpp.shape.size() < 3) { - return out; - } - - const int C = cpp.shape[0]; - const int W = cpp.shape[1]; - const int H = cpp.shape[2]; - out.resize((size_t) C * W * H); - for (int c = 0; c < C; ++c) { - for (int h = 0; h < H; ++h) { - for (int w = 0; w < W; ++w) { - out[(size_t) c * H * W + (size_t) h * W + w] = - cpp.data[(size_t) c + (size_t) w * C + (size_t) h * C * W]; - } - } - } - return out; -} - -static bool parse_phase6_case(const std::string & path, - sam3_pvs_params & params, - std::string & case_id) { - FILE * f = fopen(path.c_str(), "r"); - if (!f) { - return false; - } - - char line[1024]; - if (!fgets(line, sizeof(line), f)) { - fclose(f); - return false; - } - fclose(f); - - std::string row(line); - while (!row.empty() && (row.back() == '\n' || row.back() == '\r')) { - row.pop_back(); - } - - std::vector fields; - size_t start = 0; - while (start <= row.size()) { - size_t end = row.find('\t', start); - if (end == std::string::npos) { - end = row.size(); - } - fields.emplace_back(row.substr(start, end - start)); - start = end + 1; - if (end == row.size()) { - break; - } - } - while (fields.size() < 5) { - fields.emplace_back(); - } - - case_id = fields[0]; - params.multimask = fields[1] == "1"; - params.use_box = false; - - if (!fields[4].empty()) { - float vals[4] = {0, 0, 0, 0}; - size_t start = 0; - for (int i = 0; i < 4; ++i) { - size_t end = fields[4].find(':', start); - if (end == std::string::npos) end = fields[4].size(); - vals[i] = std::stof(fields[4].substr(start, end - start)); - start = end + 1; - } - params.box = {vals[0], vals[1], vals[2], vals[3]}; - params.use_box = true; - } - - auto parse_points = [](const std::string & field, std::vector & dst) { - size_t start = 0; - while (start < field.size()) { - size_t end = field.find('|', start); - if (end == std::string::npos) end = field.size(); - const std::string part = field.substr(start, end - start); - size_t sep = part.find(':'); - if (sep != std::string::npos) { - sam3_point p; - p.x = std::stof(part.substr(0, sep)); - p.y = std::stof(part.substr(sep + 1)); - dst.push_back(p); - } - start = end + 1; - } - }; - - parse_points(fields[2], params.pos_points); - parse_points(fields[3], params.neg_points); - return true; -} - -static metric_row compare_state_nhwc(const sam3_state & state, - const std::string & tensor_name, - const std::string & ref_path, - const std::string & dump_dir) { - metric_row row; - auto ref = load_ref_f32(ref_path); - row.tensor = tensor_name; - row.shape = shape_to_string(ref.shape); - - if (ref.data.empty()) { - row.note = "missing python ref"; - return row; - } - if (!sam3_dump_state_tensor(state, tensor_name, dump_dir + "/" + tensor_name)) { - row.note = "missing state tensor"; - return row; - } - auto cpp = load_ref_f32(dump_dir + "/" + tensor_name); - if (cpp.data.empty()) { - row.note = "missing dumped tensor"; - return row; - } - - auto transposed = ggml_dump_to_nhwc_flat(cpp); - if ((int) transposed.size() != ref.numel()) { - row.note = "numel mismatch"; - return row; - } - return compute_metrics(tensor_name, transposed.data(), ref.data.data(), ref.numel(), ref.shape); -} - -static metric_row compare_state_nchw(const sam3_state & state, - const std::string & tensor_name, - const std::string & ref_path, - const std::string & dump_dir) { - metric_row row; - auto ref = load_ref_f32(ref_path); - row.tensor = tensor_name; - row.shape = shape_to_string(ref.shape); - - if (ref.data.empty()) { - row.note = "missing python ref"; - return row; - } - if (!sam3_dump_state_tensor(state, tensor_name, dump_dir + "/" + tensor_name)) { - row.note = "missing state tensor"; - return row; - } - auto cpp = load_ref_f32(dump_dir + "/" + tensor_name); - if (cpp.data.empty()) { - row.note = "missing dumped tensor"; - return row; - } - - auto transposed = ggml_dump_to_nchw_flat(cpp); - if ((int) transposed.size() != ref.numel()) { - row.note = "numel mismatch"; - return row; - } - return compute_metrics(tensor_name, transposed.data(), ref.data.data(), ref.numel(), ref.shape); -} - -static bool save_ref_f32(const std::string & path, - const std::vector & data, - const std::vector & shape) { - FILE * shape_f = fopen((path + ".shape").c_str(), "w"); - if (!shape_f) { - return false; - } - for (size_t i = 0; i < shape.size(); ++i) { - fprintf(shape_f, "%s%d", i == 0 ? "" : ",", shape[i]); - } - fclose(shape_f); - - FILE * data_f = fopen((path + ".bin").c_str(), "wb"); - if (!data_f) { - return false; - } - const size_t n_written = fwrite(data.data(), sizeof(float), data.size(), data_f); - fclose(data_f); - return n_written == data.size(); -} - -static bool export_dump_as_nchw_ref(const std::string & ggml_dump_path, - const std::string & out_path) { - auto cpp = load_ref_f32(ggml_dump_path); - if (cpp.data.empty() || cpp.shape.size() < 3) { - return false; - } - const std::vector transposed = ggml_dump_to_nchw_flat(cpp); - const std::vector shape = {1, cpp.shape[0], cpp.shape[2], cpp.shape[1]}; - return save_ref_f32(out_path, transposed, shape); -} - -static metric_row compare_dump_direct(const std::string & name, - const std::string & got_path, - const std::string & ref_path) { - metric_row row; - auto got = load_ref_f32(got_path); - auto ref = load_ref_f32(ref_path); - row.tensor = name; - row.shape = shape_to_string(ref.shape); - - if (got.data.empty() || ref.data.empty()) { - row.note = "missing dump"; - return row; - } - if (got.numel() != ref.numel()) { - row.note = "numel mismatch"; - return row; - } - return compute_metrics(name, got.data.data(), ref.data.data(), ref.numel(), ref.shape); -} - -static backend_run run_backend(const std::string & backend_name, - bool use_gpu, - const std::string & model_path, - const std::string & ref_dir, - const sam3_pvs_params & pvs_params, - const std::string & phase6_case_id) { - backend_run run; - run.backend = backend_name; - - sam3_params params; - params.model_path = model_path; - params.n_threads = 4; - params.use_gpu = use_gpu; - - auto model = sam3_load_model(params); - if (!model) { - fprintf(stderr, "Failed to load model for backend=%s\n", backend_name.c_str()); - return run; - } - auto state = sam3_create_state(*model, params); - if (!state) { - fprintf(stderr, "Failed to create state for backend=%s\n", backend_name.c_str()); - return run; - } - - auto pre = load_ref_f32(ref_dir + "/preprocessed"); - if (pre.data.empty()) { - fprintf(stderr, "Missing preprocessed tensor in %s\n", ref_dir.c_str()); - return run; - } - - const std::string dump_root = "/tmp/sam3_pyref_" + backend_name; - ensure_dir(dump_root); - - auto t0 = std::chrono::high_resolution_clock::now(); - const bool ok = sam3_encode_image_from_preprocessed(*state, *model, pre.data.data(), 1008); - auto t1 = std::chrono::high_resolution_clock::now(); - run.encode_ms = std::chrono::duration(t1 - t0).count(); - if (!ok) { - fprintf(stderr, "sam3_encode_image_from_preprocessed failed for backend=%s\n", backend_name.c_str()); - return run; - } - - const char * nhwc_tensors[] = { - "dbg_patch_embed", - "dbg_after_pos_embed", - "dbg_after_ln_pre", - "dbg_block_0_out", - "dbg_block_1_out", - "dbg_block_2_out", - "dbg_block_10_out", - "dbg_block_14_out", - "dbg_block_15_out", - "dbg_block_16_out", - }; - for (const char * name : nhwc_tensors) { - std::string ref_name = name; - if (strncmp(name, "dbg_", 4) == 0) { - ref_name = std::string(name + 4); - } - run.encoder_rows.push_back(compare_state_nhwc(*state, name, ref_dir + "/" + ref_name, dump_root)); - } - - run.encoder_rows.push_back(compare_state_nchw(*state, "vit_output", ref_dir + "/vit_output_bchw", dump_root)); - run.encoder_rows.push_back(compare_state_nchw(*state, "neck_trk_0", ref_dir + "/neck_trk_0", dump_root)); - run.encoder_rows.push_back(compare_state_nchw(*state, "neck_trk_1", ref_dir + "/neck_trk_1", dump_root)); - run.encoder_rows.push_back(compare_state_nchw(*state, "neck_trk_2", ref_dir + "/neck_trk_2", dump_root)); - - const std::string phase6_inputs = dump_root + "/phase6_inputs"; - const std::string phase6_dump = dump_root + "/phase6_out"; - ensure_dir(phase6_inputs); - ensure_dir(phase6_dump); - - if (!export_dump_as_nchw_ref(dump_root + "/neck_trk_0", phase6_inputs + "/neck_trk_0") || - !export_dump_as_nchw_ref(dump_root + "/neck_trk_1", phase6_inputs + "/neck_trk_1") || - !export_dump_as_nchw_ref(dump_root + "/neck_trk_2", phase6_inputs + "/neck_trk_2")) { - metric_row row; - row.tensor = "phase6_bridge"; - row.note = "failed to export neck_trk_* as NCHW ref inputs"; - run.final_rows.push_back(std::move(row)); - return run; - } - - state.reset(); - sam3_free_model(*model); - model.reset(); - - sam3_params p6_params; - p6_params.model_path = model_path; - p6_params.n_threads = 4; - p6_params.use_gpu = false; - - auto phase6_model = sam3_load_model(p6_params); - if (!phase6_model) { - metric_row row; - row.tensor = "phase6_model"; - row.note = "failed to load CPU model for phase6 compare"; - run.final_rows.push_back(std::move(row)); - return run; - } - - auto t2 = std::chrono::high_resolution_clock::now(); - const bool phase6_ok = sam3_test_dump_phase6_from_ref_inputs(*phase6_model, phase6_inputs, - pvs_params, phase6_dump, - p6_params.n_threads); - auto t3 = std::chrono::high_resolution_clock::now(); - run.final_ms = std::chrono::duration(t3 - t2).count(); - if (!phase6_ok) { - metric_row row; - row.tensor = "phase6_compare"; - row.note = "sam3_test_dump_phase6_from_ref_inputs failed"; - run.final_rows.push_back(std::move(row)); - return run; - } - - const std::string p6_ref_dir = ref_dir + "/phase6/" + phase6_case_id; - const char * final_tensors[] = { - "sam_dec_masks", - "sam_dec_iou", - "sam_dec_obj_score", - "sam_dec_sam_token", - }; - for (const char * name : final_tensors) { - run.final_rows.push_back(compare_dump_direct(name, phase6_dump + "/" + name, p6_ref_dir + "/" + name)); - } - - return run; -} - -static void print_table(const std::string & title, const std::vector & rows) { - fprintf(stderr, "\n%s\n", title.c_str()); - fprintf(stderr, "| tensor | shape | max abs diff | mean abs diff | mean rel err | cosine | note |\n"); - fprintf(stderr, "|--------|-------|--------------|---------------|--------------|--------|------|\n"); - for (const auto & row : rows) { - if (!row.measured) { - fprintf(stderr, "| %s | %s | n/a | n/a | n/a | n/a | %s |\n", - row.tensor.c_str(), row.shape.c_str(), row.note.c_str()); - continue; - } - fprintf(stderr, "| %s | %s | %.6e | %.6e | %.6e | %.8f | %s |\n", - row.tensor.c_str(), row.shape.c_str(), row.max_abs, - row.mean_abs, row.mean_rel, row.cosine, row.note.c_str()); - } -} - -static void print_summary(const backend_run & run) { - fprintf(stderr, "\n=== %s vs Python ===\n", run.backend.c_str()); - fprintf(stderr, "encode_ms=%.3f final_ms=%.3f\n", run.encode_ms, run.final_ms); - print_table("Encoder / Neck", run.encoder_rows); - print_table("Final Decoder Outputs", run.final_rows); -} - -int main(int argc, char ** argv) { - if (argc < 2) { - fprintf(stderr, - "Usage: %s [ref_dir] [backend=metal|cpu|both] [phase6_case_tsv]\n", - argv[0]); - fprintf(stderr, - "Defaults: ref_dir=tests/debug_pipeline/e2e_ref backend=both phase6_case_tsv=tests/debug_pipeline/phase6_cat_box.tsv\n"); - return 1; - } - - const std::string model_path = argv[1]; - const std::string ref_dir = argc >= 3 ? argv[2] : "tests/debug_pipeline/e2e_ref"; - const std::string backend_mode = argc >= 4 ? argv[3] : "both"; - const std::string case_tsv = argc >= 5 ? argv[4] : "tests/debug_pipeline/phase6_cat_box.tsv"; - - sam3_pvs_params pvs_params; - std::string phase6_case_id; - if (!parse_phase6_case(case_tsv, pvs_params, phase6_case_id)) { - fprintf(stderr, "Failed to parse phase 6 case file: %s\n", case_tsv.c_str()); - return 1; - } - - std::vector runs; - if (backend_mode == "cpu" || backend_mode == "both") { - runs.push_back(run_backend("cpu", false, model_path, ref_dir, pvs_params, phase6_case_id)); - } - if (backend_mode == "metal" || backend_mode == "both") { - runs.push_back(run_backend("metal", true, model_path, ref_dir, pvs_params, phase6_case_id)); - } - - for (const auto & run : runs) { - print_summary(run); - } - - return 0; -} diff --git a/tests/test_metal_real_profile.cpp b/tests/test_metal_real_profile.cpp deleted file mode 100644 index c38d8b1..0000000 --- a/tests/test_metal_real_profile.cpp +++ /dev/null @@ -1,108 +0,0 @@ -// Profile the REAL SAM3 ViT graph on Metal. -// Runs encode_image twice to separate first-run compilation from steady-state. -// Also breaks down graph_compute into encoding vs GPU execution. - -#include "sam3.h" -#include "ggml-backend.h" - -#ifdef GGML_USE_METAL -#include "ggml-metal.h" -#endif - -#include -#include -#include - -static double now_ms() { - return std::chrono::duration( - std::chrono::high_resolution_clock::now().time_since_epoch()).count(); -} - -int main(int argc, char** argv) { - if (argc < 3) { - fprintf(stderr, "Usage: %s \n", argv[0]); - return 1; - } - - auto img = sam3_load_image(argv[2]); - fprintf(stderr, "Image: %dx%d\n\n", img.width, img.height); - - // ── Metal: run encode_image TWICE ─────────────────────────────────── - fprintf(stderr, "=== Metal: Loading model ===\n"); - sam3_params params; - params.model_path = argv[1]; - params.use_gpu = true; - params.n_threads = 8; - - double t0 = now_ms(); - auto model = sam3_load_model(params); - double t1 = now_ms(); - fprintf(stderr, " Model load: %.0f ms\n", t1 - t0); - - if (!model) { fprintf(stderr, "Failed to load model\n"); return 1; } - - auto state = sam3_create_state(*model, params); - if (!state) { fprintf(stderr, "Failed to create state\n"); return 1; } - - // Run 1: includes pipeline compilation - fprintf(stderr, "\n=== Metal: encode_image RUN 1 (cold — includes pipeline compilation) ===\n"); - t0 = now_ms(); - bool ok = sam3_encode_image(*state, *model, img); - t1 = now_ms(); - fprintf(stderr, " Run 1 total: %.0f ms %s\n", t1 - t0, ok ? "" : "FAILED"); - - // Read back a tensor to force GPU synchronization - { - double ts0 = now_ms(); - sam3_dump_state_tensor(*state, "vit_output", "/tmp/sam3_metal_prof_vit"); - double ts1 = now_ms(); - fprintf(stderr, " Readback sync: %.0f ms\n", ts1 - ts0); - } - - // Run 2: pipelines already compiled, GPU warmed up - fprintf(stderr, "\n=== Metal: encode_image RUN 2 (warm — no compilation) ===\n"); - t0 = now_ms(); - ok = sam3_encode_image(*state, *model, img); - t1 = now_ms(); - fprintf(stderr, " Run 2 total: %.0f ms %s\n", t1 - t0, ok ? "" : "FAILED"); - - { - double ts0 = now_ms(); - sam3_dump_state_tensor(*state, "vit_output", "/tmp/sam3_metal_prof_vit2"); - double ts1 = now_ms(); - fprintf(stderr, " Readback sync: %.0f ms\n", ts1 - ts0); - } - - // Run 3: one more to confirm steady state - fprintf(stderr, "\n=== Metal: encode_image RUN 3 (warm) ===\n"); - t0 = now_ms(); - ok = sam3_encode_image(*state, *model, img); - t1 = now_ms(); - fprintf(stderr, " Run 3 total: %.0f ms %s\n", t1 - t0, ok ? "" : "FAILED"); - - { - double ts0 = now_ms(); - sam3_dump_state_tensor(*state, "vit_output", "/tmp/sam3_metal_prof_vit3"); - double ts1 = now_ms(); - fprintf(stderr, " Readback sync: %.0f ms\n", ts1 - ts0); - } - - // ── CPU: for comparison ───────────────────────────────────────────── - fprintf(stderr, "\n=== CPU: encode_image (for comparison) ===\n"); - { - sam3_params cpu_params; - cpu_params.model_path = argv[1]; - cpu_params.use_gpu = false; - cpu_params.n_threads = 8; - - auto cpu_model = sam3_load_model(cpu_params); - auto cpu_state = sam3_create_state(*cpu_model, cpu_params); - - t0 = now_ms(); - ok = sam3_encode_image(*cpu_state, *cpu_model, img); - t1 = now_ms(); - fprintf(stderr, " CPU total: %.0f ms %s\n", t1 - t0, ok ? "" : "FAILED"); - } - - return 0; -} diff --git a/tests/test_metal_vit.cpp b/tests/test_metal_vit.cpp deleted file mode 100644 index 9b5fef6..0000000 --- a/tests/test_metal_vit.cpp +++ /dev/null @@ -1,166 +0,0 @@ -// Test Metal ViT encoding: compare CPU vs Metal image encoder outputs. -// Usage: test_metal_vit - -#include "sam3.h" - -#include "ggml-backend.h" -#include "ggml.h" - -#include -#include -#include -#include -#include -#include - -int main(int argc, char** argv) { - if (argc < 3) { - fprintf(stderr, "Usage: %s \n", argv[0]); - return 1; - } - - const std::string model_path = argv[1]; - const std::string image_path = argv[2]; - - // Load image once - auto img = sam3_load_image(image_path); - if (img.data.empty()) { - fprintf(stderr, "Failed to load image: %s\n", image_path.c_str()); - return 1; - } - fprintf(stderr, "Loaded %dx%d image\n", img.width, img.height); - - // ── Run CPU encoding ──────────────────────────────────────────────── - fprintf(stderr, "\n=== CPU Encoding ===\n"); - std::vector cpu_neck0, cpu_neck1, cpu_neck2; - { - sam3_params params; - params.model_path = model_path; - params.use_gpu = false; - params.n_threads = 8; - - auto model = sam3_load_model(params); - if (!model) { fprintf(stderr, "Failed to load model (CPU)\n"); return 1; } - - auto state = sam3_create_state(*model, params); - if (!state) { fprintf(stderr, "Failed to create state (CPU)\n"); return 1; } - - auto t0 = std::chrono::high_resolution_clock::now(); - bool ok = sam3_encode_image(*state, *model, img); - auto t1 = std::chrono::high_resolution_clock::now(); - double ms = std::chrono::duration(t1 - t0).count(); - - if (!ok) { fprintf(stderr, "CPU encode_image failed\n"); return 1; } - fprintf(stderr, " CPU encode_image: %.1f ms\n", ms); - - // Dump tensors via the debug interface - sam3_dump_state_tensor(*state, "vit_output", "/tmp/sam3_cpu_vit"); - sam3_dump_state_tensor(*state, "neck_trk_0", "/tmp/sam3_cpu_neck0"); - sam3_dump_state_tensor(*state, "neck_trk_1", "/tmp/sam3_cpu_neck1"); - sam3_dump_state_tensor(*state, "neck_trk_2", "/tmp/sam3_cpu_neck2"); - } - - // ── Run Metal encoding ────────────────────────────────────────────── - fprintf(stderr, "\n=== Metal Encoding ===\n"); - { - sam3_params params; - params.model_path = model_path; - params.use_gpu = true; - params.n_threads = 8; - - auto model = sam3_load_model(params); - if (!model) { fprintf(stderr, "Failed to load model (Metal)\n"); return 1; } - - auto state = sam3_create_state(*model, params); - if (!state) { fprintf(stderr, "Failed to create state (Metal)\n"); return 1; } - - auto t0 = std::chrono::high_resolution_clock::now(); - bool ok = sam3_encode_image(*state, *model, img); - auto t1 = std::chrono::high_resolution_clock::now(); - double ms = std::chrono::duration(t1 - t0).count(); - - if (!ok) { fprintf(stderr, "Metal encode_image failed\n"); return 1; } - fprintf(stderr, " Metal encode_image: %.1f ms\n", ms); - - sam3_dump_state_tensor(*state, "vit_output", "/tmp/sam3_metal_vit"); - sam3_dump_state_tensor(*state, "neck_trk_0", "/tmp/sam3_metal_neck0"); - sam3_dump_state_tensor(*state, "neck_trk_1", "/tmp/sam3_metal_neck1"); - sam3_dump_state_tensor(*state, "neck_trk_2", "/tmp/sam3_metal_neck2"); - } - - // ── Compare outputs ───────────────────────────────────────────────── - fprintf(stderr, "\n=== Comparing CPU vs Metal ===\n"); - bool all_pass = true; - - auto compare_dumps = [&](const std::string& name, - const std::string& cpu_path, - const std::string& metal_path, - float atol) { - // Load shapes - auto load_bin = [](const std::string& path, std::vector& shape) -> std::vector { - // Read shape - { - FILE* f = fopen((path + ".shape").c_str(), "r"); - if (!f) return {}; - char buf[256]; - if (fgets(buf, sizeof(buf), f)) { - char* tok = strtok(buf, ",\n"); - while (tok) { - shape.push_back(atoi(tok)); - tok = strtok(nullptr, ",\n"); - } - } - fclose(f); - } - // Read data - int n = 1; - for (auto d : shape) n *= d; - std::vector data(n); - FILE* f = fopen((path + ".bin").c_str(), "rb"); - if (!f) return {}; - fread(data.data(), sizeof(float), n, f); - fclose(f); - return data; - }; - - std::vector cpu_shape, metal_shape; - auto cpu_data = load_bin(cpu_path, cpu_shape); - auto metal_data = load_bin(metal_path, metal_shape); - - if (cpu_data.empty() || metal_data.empty()) { - fprintf(stderr, " [SKIP] %s — could not load dumps\n", name.c_str()); - return; - } - if (cpu_data.size() != metal_data.size()) { - fprintf(stderr, " [FAIL] %s — size mismatch: %zu vs %zu\n", - name.c_str(), cpu_data.size(), metal_data.size()); - all_pass = false; - return; - } - - float max_diff = 0.0f; - int max_idx = 0; - int n_bad = 0; - for (size_t i = 0; i < cpu_data.size(); i++) { - float diff = std::fabs(cpu_data[i] - metal_data[i]); - if (diff > max_diff) { max_diff = diff; max_idx = (int)i; } - if (diff > atol) n_bad++; - } - - bool ok = (max_diff < atol); - fprintf(stderr, " %s %s: max_diff=%.6f at [%d] (cpu=%.6f metal=%.6f) n_bad=%d/%zu\n", - ok ? "[PASS]" : "[FAIL]", name.c_str(), max_diff, max_idx, - cpu_data[max_idx], metal_data[max_idx], n_bad, cpu_data.size()); - if (!ok) all_pass = false; - }; - - // ViT backbone + FPN neck outputs — expect reasonable tolerance - // (different execution order on GPU can cause FP32 rounding differences) - compare_dumps("vit_output", "/tmp/sam3_cpu_vit", "/tmp/sam3_metal_vit", 0.1f); - compare_dumps("neck_trk_0", "/tmp/sam3_cpu_neck0", "/tmp/sam3_metal_neck0", 0.1f); - compare_dumps("neck_trk_1", "/tmp/sam3_cpu_neck1", "/tmp/sam3_metal_neck1", 0.1f); - compare_dumps("neck_trk_2", "/tmp/sam3_cpu_neck2", "/tmp/sam3_metal_neck2", 0.1f); - - fprintf(stderr, "\n%s\n", all_pass ? "ALL COMPARISONS PASSED" : "SOME COMPARISONS FAILED"); - return all_pass ? 0 : 1; -} diff --git a/tests/test_metal_vit_prefix_compare.cpp b/tests/test_metal_vit_prefix_compare.cpp deleted file mode 100644 index 8f75a9a..0000000 --- a/tests/test_metal_vit_prefix_compare.cpp +++ /dev/null @@ -1,187 +0,0 @@ -#include "sam3.h" -#include "test_utils.h" - -#include -#include -#include -#include - -static std::string format_shape(const std::vector & shape) { - std::string out = "["; - for (size_t i = 0; i < shape.size(); ++i) { - if (i > 0) { - out += ", "; - } - out += std::to_string(shape[i]); - } - out += "]"; - return out; -} - -static std::vector trim_shape(const int64_t ne[4]) { - std::vector shape = { - (int) ne[0], - (int) ne[1], - (int) ne[2], - (int) ne[3], - }; - while (!shape.empty() && shape.back() == 1) { - shape.pop_back(); - } - if (shape.empty()) { - shape.push_back(1); - } - return shape; -} - -static ref_tensor_f32 run_prefix_stage_or_die(const sam3_model & model, - sam3_vit_prefix_stage stage, - const ref_tensor_f32 & input, - int n_threads) { - int64_t in_ne[4] = {1, 1, 1, 1}; - for (size_t i = 0; i < input.shape.size() && i < 4; ++i) { - in_ne[i] = input.shape[i]; - } - - int64_t out_ne[4] = {0, 0, 0, 0}; - ref_tensor_f32 out; - if (!sam3_test_run_vit_prefix_stage(model, stage, input.data.data(), in_ne, out.data, out_ne, n_threads)) { - fprintf(stderr, "prefix stage %d failed\n", (int) stage); - std::exit(1); - } - out.shape = trim_shape(out_ne); - return out; -} - -static ref_tensor_f32 make_image_input(const ref_tensor_f32 & preprocessed, int img_size) { - ref_tensor_f32 image = preprocessed; - image.shape = {img_size, img_size, 3, 1}; - return image; -} - -static ref_tensor_f32 run_block0_input_or_die(const sam3_model & model, - const ref_tensor_f32 & image, - int img_size, - int n_threads) { - int64_t out_ne[4] = {0, 0, 0, 0}; - ref_tensor_f32 out; - if (!sam3_test_run_vit_block0_input(model, image.data.data(), img_size, out.data, out_ne, n_threads)) { - fprintf(stderr, "block0 input prefix failed\n"); - std::exit(1); - } - out.shape = trim_shape(out_ne); - return out; -} - -static void print_diff_row(const char * label, - const ref_tensor_f32 & a, - const ref_tensor_f32 & b, - float atol = 1e-4f) { - const compare_result r = compare_tensors(a.data.data(), b.data.data(), a.numel(), atol); - fprintf(stderr, "%-18s %-18s %12.6f %12.6f %10d\n", - label, - format_shape(a.shape).c_str(), - r.max_diff, - r.mean_diff, - r.n_bad); -} - -int main(int argc, char ** argv) { - if (argc < 2) { - fprintf(stderr, "Usage: %s [ref_dir]\n", argv[0]); - fprintf(stderr, "Default ref_dir: tests/ref_phase3\n"); - return 1; - } - - const std::string model_path = argv[1]; - const std::string ref_dir = argc >= 3 ? argv[2] : "tests/ref_phase3"; - - const ref_tensor_f32 preprocessed = load_ref_f32(ref_dir + "/preprocessed"); - if (preprocessed.data.empty() || preprocessed.shape.size() != 4) { - fprintf(stderr, "failed to load %s/preprocessed\n", ref_dir.c_str()); - return 1; - } - - const int img_size = preprocessed.shape[2]; - const int n_threads = 8; - - sam3_params cpu_params; - cpu_params.model_path = model_path; - cpu_params.use_gpu = false; - cpu_params.n_threads = n_threads; - - sam3_params metal_params = cpu_params; - metal_params.use_gpu = true; - - auto cpu_model = sam3_load_model(cpu_params); - auto metal_model = sam3_load_model(metal_params); - if (!cpu_model || !metal_model) { - fprintf(stderr, "failed to load CPU or Metal model\n"); - return 1; - } - - fprintf(stderr, "\n=== Exact ViT prefix compare ===\n"); - fprintf(stderr, "%-18s %-18s %12s %12s %10s\n", - "stage", "shape", "max_abs_diff", "mean_abs", "n_bad"); - - const ref_tensor_f32 image_input = make_image_input(preprocessed, img_size); - - const ref_tensor_f32 cpu_im2col = run_prefix_stage_or_die(*cpu_model, SAM3_VIT_PREFIX_STAGE_PATCH_IM2COL, image_input, n_threads); - const ref_tensor_f32 metal_im2col = run_prefix_stage_or_die(*metal_model, SAM3_VIT_PREFIX_STAGE_PATCH_IM2COL, image_input, n_threads); - print_diff_row("patch_im2col", cpu_im2col, metal_im2col); - - const ref_tensor_f32 cpu_patch_raw_shared = run_prefix_stage_or_die(*cpu_model, SAM3_VIT_PREFIX_STAGE_PATCH_MULMAT_RAW, cpu_im2col, n_threads); - const ref_tensor_f32 metal_patch_raw_shared = run_prefix_stage_or_die(*metal_model, SAM3_VIT_PREFIX_STAGE_PATCH_MULMAT_RAW, cpu_im2col, n_threads); - print_diff_row("patch_mulmat_raw", cpu_patch_raw_shared, metal_patch_raw_shared); - - { - ref_tensor_f32 host_patch_raw_f32; - ref_tensor_f32 host_patch_raw_f64; - int64_t out_ne_f32[4] = {0, 0, 0, 0}; - int64_t out_ne_f64[4] = {0, 0, 0, 0}; - int64_t in_ne[4] = {1, 1, 1, 1}; - for (size_t i = 0; i < cpu_im2col.shape.size() && i < 4; ++i) { - in_ne[i] = cpu_im2col.shape[i]; - } - - if (!sam3_test_run_patch_mulmat_host_ref(*cpu_model, cpu_im2col.data.data(), in_ne, false, host_patch_raw_f32.data, out_ne_f32) || - !sam3_test_run_patch_mulmat_host_ref(*cpu_model, cpu_im2col.data.data(), in_ne, true, host_patch_raw_f64.data, out_ne_f64)) { - fprintf(stderr, "host patch_mulmat reference failed\n"); - return 1; - } - - host_patch_raw_f32.shape = trim_shape(out_ne_f32); - host_patch_raw_f64.shape = trim_shape(out_ne_f64); - - print_diff_row("patch_raw_cpu_f32", cpu_patch_raw_shared, host_patch_raw_f32); - print_diff_row("patch_raw_cpu_f64", cpu_patch_raw_shared, host_patch_raw_f64); - print_diff_row("patch_raw_mtl_f32", metal_patch_raw_shared, host_patch_raw_f32); - print_diff_row("patch_raw_mtl_f64", metal_patch_raw_shared, host_patch_raw_f64); - } - - const ref_tensor_f32 cpu_patch_shared = run_prefix_stage_or_die(*cpu_model, SAM3_VIT_PREFIX_STAGE_PATCH_MULMAT, cpu_im2col, n_threads); - const ref_tensor_f32 metal_patch_shared = run_prefix_stage_or_die(*metal_model, SAM3_VIT_PREFIX_STAGE_PATCH_MULMAT, cpu_im2col, n_threads); - print_diff_row("patch_mulmat", cpu_patch_shared, metal_patch_shared); - - const ref_tensor_f32 cpu_patch = run_prefix_stage_or_die(*cpu_model, SAM3_VIT_PREFIX_STAGE_PATCH_EMBED, image_input, n_threads); - const ref_tensor_f32 metal_patch = run_prefix_stage_or_die(*metal_model, SAM3_VIT_PREFIX_STAGE_PATCH_EMBED, image_input, n_threads); - print_diff_row("patch_embed", cpu_patch, metal_patch); - - const ref_tensor_f32 cpu_pos_shared = run_prefix_stage_or_die(*cpu_model, SAM3_VIT_PREFIX_STAGE_POS_ADD, cpu_patch, n_threads); - const ref_tensor_f32 metal_pos_shared = run_prefix_stage_or_die(*metal_model, SAM3_VIT_PREFIX_STAGE_POS_ADD, cpu_patch, n_threads); - print_diff_row("after_pos_embed", cpu_pos_shared, metal_pos_shared); - - const ref_tensor_f32 cpu_ln_norm_shared = run_prefix_stage_or_die(*cpu_model, SAM3_VIT_PREFIX_STAGE_LN_PRE_NORM, cpu_pos_shared, n_threads); - const ref_tensor_f32 metal_ln_norm_shared = run_prefix_stage_or_die(*metal_model, SAM3_VIT_PREFIX_STAGE_LN_PRE_NORM, cpu_pos_shared, n_threads); - print_diff_row("ln_pre_norm", cpu_ln_norm_shared, metal_ln_norm_shared); - - const ref_tensor_f32 cpu_ln_pre_shared = run_prefix_stage_or_die(*cpu_model, SAM3_VIT_PREFIX_STAGE_LN_PRE, cpu_pos_shared, n_threads); - const ref_tensor_f32 metal_ln_pre_shared = run_prefix_stage_or_die(*metal_model, SAM3_VIT_PREFIX_STAGE_LN_PRE, cpu_pos_shared, n_threads); - print_diff_row("after_ln_pre", cpu_ln_pre_shared, metal_ln_pre_shared); - - const ref_tensor_f32 cpu_block0_input = run_block0_input_or_die(*cpu_model, preprocessed, img_size, n_threads); - const ref_tensor_f32 metal_block0_input = run_block0_input_or_die(*metal_model, preprocessed, img_size, n_threads); - print_diff_row("block0_input", cpu_block0_input, metal_block0_input); - - return 0; -} diff --git a/tests/test_metal_vit_upstream_compare.cpp b/tests/test_metal_vit_upstream_compare.cpp deleted file mode 100644 index 420115c..0000000 --- a/tests/test_metal_vit_upstream_compare.cpp +++ /dev/null @@ -1,383 +0,0 @@ -#include "ggml-alloc.h" -#include "ggml-backend.h" -#include "ggml-cpu.h" -#include "ggml.h" - -#ifdef GGML_USE_METAL -#include "ggml-metal.h" -#endif - -#include "sam3.h" -#include "test_utils.h" - -#include -#include -#include -#include -#include -#include - -struct tensor_spec { - std::string label; - std::string name; -}; - -struct backend_run_result { - double elapsed_ms = 0.0; - std::map infos; -}; - -struct isolated_run_result { - double elapsed_ms = 0.0; - std::vector output; - std::vector shape; -}; - -static std::string format_shape(const std::vector & shape) { - std::string out = "["; - for (size_t i = 0; i < shape.size(); ++i) { - if (i > 0) { - out += ", "; - } - out += std::to_string(shape[i]); - } - out += "]"; - return out; -} - -static std::string make_block_name(int block_idx, const char * suffix) { - char name[64]; - snprintf(name, sizeof(name), "dbg_block_%d_%s", block_idx, suffix); - return name; -} - -static bool dump_named_tensors(const sam3_state & state, - const std::vector & names, - const std::string & output_dir, - std::map & infos) { - for (const auto & name : names) { - sam3_tensor_info info; - if (!sam3_get_state_tensor_info(state, name, info)) { - fprintf(stderr, "missing tensor info for %s\n", name.c_str()); - return false; - } - infos[name] = info; - if (!sam3_dump_state_tensor(state, name, output_dir + "/" + name)) { - fprintf(stderr, "failed to dump %s\n", name.c_str()); - return false; - } - } - return true; -} - -static ggml_backend_t create_backend(bool use_gpu, int n_threads) { -#ifdef GGML_USE_METAL - if (use_gpu) { - return ggml_backend_metal_init(); - } -#else - (void) use_gpu; -#endif - ggml_backend_t backend = ggml_backend_cpu_init(); - if (backend) { - ggml_backend_cpu_set_n_threads(backend, n_threads); - } - return backend; -} - -static std::vector normalize_shape_4d(const std::vector & shape) { - std::vector out = shape; - while (out.size() < 4) { - out.push_back(1); - } - return out; -} - -static ggml_tensor * new_tensor_4d_from_shape(ggml_context * ctx, - ggml_type type, - const std::vector & shape) { - const std::vector s = normalize_shape_4d(shape); - return ggml_new_tensor_4d(ctx, type, s[0], s[1], s[2], s[3]); -} - -static bool load_block_norm_params(const std::string & model_path, - int block_idx, - ref_tensor_f32 & weight, - ref_tensor_f32 & bias) { - sam3_params params; - params.model_path = model_path; - params.use_gpu = false; - params.n_threads = 1; - - auto model = sam3_load_model(params); - if (!model) { - return false; - } - - const std::string tmp_dir = "/tmp/sam3_vit_upstream_model"; - ensure_dir(tmp_dir); - const std::string prefix = "vit.blocks." + std::to_string(block_idx) + ".norm1"; - if (!sam3_dump_model_tensor(*model, prefix + ".weight", tmp_dir + "/norm1_weight")) { - return false; - } - if (!sam3_dump_model_tensor(*model, prefix + ".bias", tmp_dir + "/norm1_bias")) { - return false; - } - - weight = load_ref_f32(tmp_dir + "/norm1_weight"); - bias = load_ref_f32(tmp_dir + "/norm1_bias"); - sam3_free_model(*model); - model.reset(); - return !weight.data.empty() && !bias.data.empty(); -} - -static isolated_run_result run_isolated_norm(ggml_backend_t backend, - const ref_tensor_f32 & x_in, - const ref_tensor_f32 & weight, - const ref_tensor_f32 & bias) { - isolated_run_result result; - - const size_t ctx_size = ggml_tensor_overhead() * 16 + ggml_graph_overhead(); - ggml_init_params params = { - /*.mem_size =*/ ctx_size, - /*.mem_buffer =*/ nullptr, - /*.no_alloc =*/ true, - }; - ggml_context * ctx = ggml_init(params); - if (!ctx) { - return result; - } - - ggml_tensor * x = new_tensor_4d_from_shape(ctx, GGML_TYPE_F32, x_in.shape); - ggml_tensor * w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, weight.shape[0]); - ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, bias.shape[0]); - ggml_set_input(x); - ggml_set_input(w); - ggml_set_input(b); - - ggml_tensor * y = ggml_norm(ctx, x, 1e-5f); - y = ggml_mul_inplace(ctx, y, w); - y = ggml_add_inplace(ctx, y, b); - ggml_set_output(y); - - ggml_cgraph * graph = ggml_new_graph(ctx); - ggml_build_forward_expand(graph, y); - - ggml_gallocr_t galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!ggml_gallocr_reserve(galloc, graph) || !ggml_gallocr_alloc_graph(galloc, graph)) { - ggml_gallocr_free(galloc); - ggml_free(ctx); - return result; - } - - ggml_backend_tensor_set(x, x_in.data.data(), 0, x_in.data.size() * sizeof(float)); - ggml_backend_tensor_set(w, weight.data.data(), 0, weight.data.size() * sizeof(float)); - ggml_backend_tensor_set(b, bias.data.data(), 0, bias.data.size() * sizeof(float)); - - auto t0 = std::chrono::high_resolution_clock::now(); - const ggml_status status = ggml_backend_graph_compute(backend, graph); - auto t1 = std::chrono::high_resolution_clock::now(); - result.elapsed_ms = std::chrono::duration(t1 - t0).count(); - if (status == GGML_STATUS_SUCCESS) { - result.shape = { - (int) y->ne[0], - (int) y->ne[1], - (int) y->ne[2], - (int) y->ne[3], - }; - result.output.resize((size_t) ggml_nelements(y)); - ggml_backend_tensor_get(y, result.output.data(), 0, result.output.size() * sizeof(float)); - } else { - fprintf(stderr, "isolated norm failed: %s\n", ggml_status_to_string(status)); - } - - ggml_gallocr_free(galloc); - ggml_free(ctx); - return result; -} - -static bool run_backend(const std::string & model_path, - const std::vector & chw_data, - int img_size, - bool use_gpu, - int n_threads, - int debug_block_idx, - const std::vector & names, - const std::string & dump_dir, - backend_run_result & result) { - sam3_params params; - params.model_path = model_path; - params.use_gpu = use_gpu; - params.n_threads = n_threads; - - auto model = sam3_load_model(params); - if (!model) { - fprintf(stderr, "failed to load model for %s\n", use_gpu ? "Metal" : "CPU"); - return false; - } - - auto state = sam3_create_state(*model, params); - if (!state) { - fprintf(stderr, "failed to create state for %s\n", use_gpu ? "Metal" : "CPU"); - return false; - } - - auto t0 = std::chrono::high_resolution_clock::now(); - const bool ok = sam3_encode_vit_from_preprocessed_selective( - *state, *model, chw_data.data(), img_size, names); - auto t1 = std::chrono::high_resolution_clock::now(); - result.elapsed_ms = std::chrono::duration(t1 - t0).count(); - if (!ok) { - fprintf(stderr, "%s selective ViT encode failed\n", use_gpu ? "Metal" : "CPU"); - return false; - } - - if (!dump_named_tensors(*state, names, dump_dir, result.infos)) { - return false; - } - - state.reset(); - sam3_free_model(*model); - model.reset(); - return true; -} - -static std::vector default_specs_for_block(int block_idx) { - std::vector specs; - - if (block_idx <= 0) { - specs.push_back({"block_0_input / norm1_input", "dbg_after_ln_pre"}); - } else { - specs.push_back({"block_input / norm1_input", make_block_name(block_idx - 1, "out")}); - } - - specs.push_back({"norm1_output / qkv_input", make_block_name(block_idx, "norm1")}); - specs.push_back({"qkv_proj_output", make_block_name(block_idx, "qkv_proj")}); - specs.push_back({"attn_output", make_block_name(block_idx, "attn_out")}); - specs.push_back({"attn_proj_output", make_block_name(block_idx, "attn_proj")}); - specs.push_back({"resid1_output", make_block_name(block_idx, "resid1")}); - specs.push_back({"mlp_output", make_block_name(block_idx, "mlp")}); - specs.push_back({"block_output", make_block_name(block_idx, "out")}); - specs.push_back({"vit_output", "vit_output"}); - - return specs; -} - -int main(int argc, char ** argv) { - if (argc < 2) { - fprintf(stderr, "Usage: %s [ref_dir] [block_idx]\n", argv[0]); - fprintf(stderr, "Default ref_dir: tests/ref_phase3\n"); - fprintf(stderr, "Default block_idx: 15\n"); - return 1; - } - - const std::string model_path = argv[1]; - const std::string ref_dir = argc >= 3 ? argv[2] : "tests/ref_phase3"; - const int block_idx = argc >= 4 ? std::atoi(argv[3]) : 15; - if (block_idx < 0 || block_idx >= 32) { - fprintf(stderr, "invalid block_idx=%d\n", block_idx); - return 1; - } - - auto preprocessed = load_ref_f32(ref_dir + "/preprocessed"); - if (preprocessed.data.empty() || preprocessed.shape.size() != 4) { - fprintf(stderr, "failed to load %s/preprocessed\n", ref_dir.c_str()); - return 1; - } - const int img_size = preprocessed.shape[2]; - - const std::vector specs = default_specs_for_block(block_idx); - std::vector names; - names.reserve(specs.size()); - for (const auto & spec : specs) { - names.push_back(spec.name); - } - - const std::string cpu_dir = "/tmp/sam3_cpu_vit_upstream"; - const std::string metal_dir = "/tmp/sam3_metal_vit_upstream"; - ensure_dir(cpu_dir); - ensure_dir(metal_dir); - - backend_run_result cpu_run; - backend_run_result metal_run; - - fprintf(stderr, "\n=== CPU block %d selective ViT run ===\n", block_idx); - if (!run_backend(model_path, preprocessed.data, img_size, false, 8, block_idx, names, cpu_dir, cpu_run)) { - return 1; - } - fprintf(stderr, "CPU wall time: %.1f ms\n", cpu_run.elapsed_ms); - - fprintf(stderr, "\n=== Metal block %d selective ViT run ===\n", block_idx); - if (!run_backend(model_path, preprocessed.data, img_size, true, 8, block_idx, names, metal_dir, metal_run)) { - return 1; - } - fprintf(stderr, "Metal wall time: %.1f ms\n", metal_run.elapsed_ms); - - fprintf(stderr, "\n=== CPU vs Metal block %d ===\n", block_idx); - fprintf(stderr, "%-28s %-24s %-18s %14s %14s %12s\n", - "semantic", "tensor", "shape", "max_abs_diff", "mean_abs_diff", "n_bad"); - - for (const auto & spec : specs) { - auto cpu = load_ref_f32(cpu_dir + "/" + spec.name); - auto metal = load_ref_f32(metal_dir + "/" + spec.name); - if (cpu.data.empty() || metal.data.empty()) { - fprintf(stderr, "%-28s %-24s %-18s %14s %14s %12s\n", - spec.label.c_str(), spec.name.c_str(), "-", "load-fail", "load-fail", "-"); - continue; - } - - const compare_result diff = compare_tensors(cpu.data.data(), metal.data.data(), cpu.numel(), 1e-4f); - fprintf(stderr, "%-28s %-24s %-18s %14.6f %14.6f %12d\n", - spec.label.c_str(), - spec.name.c_str(), - format_shape(cpu.shape).c_str(), - diff.max_diff, - diff.mean_diff, - diff.n_bad); - } - - ref_tensor_f32 norm_w; - ref_tensor_f32 norm_b; - if (!load_block_norm_params(model_path, block_idx, norm_w, norm_b)) { - fprintf(stderr, "\nfailed to load norm1 params for block %d\n", block_idx); - return 1; - } - - const ref_tensor_f32 block_input = load_ref_f32(cpu_dir + "/" + specs[0].name); - const ref_tensor_f32 cpu_norm_dump = load_ref_f32(cpu_dir + "/" + make_block_name(block_idx, "norm1")); - const ref_tensor_f32 metal_norm_dump = load_ref_f32(metal_dir + "/" + make_block_name(block_idx, "norm1")); - if (!block_input.data.empty() && !cpu_norm_dump.data.empty() && !metal_norm_dump.data.empty()) { - ggml_backend_t cpu_backend = create_backend(false, 8); - ggml_backend_t metal_backend = create_backend(true, 8); - - const isolated_run_result iso_cpu = run_isolated_norm(cpu_backend, block_input, norm_w, norm_b); - const isolated_run_result iso_metal = run_isolated_norm(metal_backend, block_input, norm_w, norm_b); - - if (cpu_backend) { - ggml_backend_free(cpu_backend); - } - if (metal_backend) { - ggml_backend_free(metal_backend); - } - - if (!iso_cpu.output.empty() && !iso_metal.output.empty()) { - const compare_result metal_vs_cpu = compare_tensors( - iso_cpu.output.data(), iso_metal.output.data(), (int) iso_cpu.output.size(), 1e-4f); - const compare_result cpu_vs_dump = compare_tensors( - cpu_norm_dump.data.data(), iso_cpu.output.data(), cpu_norm_dump.numel(), 1e-4f); - const compare_result metal_vs_dump = compare_tensors( - metal_norm_dump.data.data(), iso_metal.output.data(), metal_norm_dump.numel(), 1e-4f); - - fprintf(stderr, "\n=== Isolated block %d norm1 ===\n", block_idx); - fprintf(stderr, "input shape: %s\n", format_shape(block_input.shape).c_str()); - fprintf(stderr, "isolated metal_vs_cpu: max_abs_diff=%.6f mean_abs_diff=%.6f n_bad=%d\n", - metal_vs_cpu.max_diff, metal_vs_cpu.mean_diff, metal_vs_cpu.n_bad); - fprintf(stderr, "isolated cpu_vs_graph_dump: max_abs_diff=%.6f mean_abs_diff=%.6f n_bad=%d\n", - cpu_vs_dump.max_diff, cpu_vs_dump.mean_diff, cpu_vs_dump.n_bad); - fprintf(stderr, "isolated metal_vs_graph_dump: max_abs_diff=%.6f mean_abs_diff=%.6f n_bad=%d\n", - metal_vs_dump.max_diff, metal_vs_dump.mean_diff, metal_vs_dump.n_bad); - } - } - - return 0; -} diff --git a/tests/test_metal_win.cpp b/tests/test_metal_win.cpp deleted file mode 100644 index 9fd296e..0000000 --- a/tests/test_metal_win.cpp +++ /dev/null @@ -1,252 +0,0 @@ -// Test Metal kernels for ggml_win_part and ggml_win_unpart. -// Runs the same graph on CPU and Metal, then compares results element-by-element. - -#include "ggml-alloc.h" -#include "ggml-backend.h" -#include "ggml-cpu.h" -#include "ggml.h" - -#ifdef GGML_USE_METAL -#include "ggml-metal.h" -#endif - -#include -#include -#include -#include -#include - -// Run a graph: build win_part -> win_unpart on the given backend. -// Returns the final output as a flat float vector. -static std::vector run_win_roundtrip(ggml_backend_t backend, - const std::vector& input, - int C, int W, int H, int WS) { - // Graph context - size_t ctx_size = ggml_tensor_overhead() * 16 + ggml_graph_overhead(); - struct ggml_init_params params = {ctx_size, nullptr, true}; - struct ggml_context* ctx = ggml_init(params); - - // Input tensor: [C, W, H, 1] - auto* inp = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, C, W, H); - ggml_set_name(inp, "input"); - ggml_set_input(inp); - - // win_part: [C, W, H, 1] -> [C, WS, WS, np] - auto* parted = ggml_win_part(ctx, inp, WS); - ggml_set_name(parted, "parted"); - - // win_unpart: [C, WS, WS, np] -> [C, W, H, 1] - auto* unparted = ggml_win_unpart(ctx, parted, W, H, WS); - ggml_set_name(unparted, "unparted"); - ggml_set_output(unparted); - - // Build graph - auto* graph = ggml_new_graph(ctx); - ggml_build_forward_expand(graph, unparted); - - // Allocate - auto* galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!ggml_gallocr_reserve(galloc, graph) || !ggml_gallocr_alloc_graph(galloc, graph)) { - fprintf(stderr, "FAIL: gallocr_alloc_graph failed\n"); - exit(1); - } - - // Set input data - ggml_backend_tensor_set(inp, input.data(), 0, input.size() * sizeof(float)); - - // Compute - ggml_backend_graph_compute(backend, graph); - - // Read output - std::vector output(C * W * H); - ggml_backend_tensor_get(unparted, output.data(), 0, output.size() * sizeof(float)); - - ggml_gallocr_free(galloc); - ggml_free(ctx); - return output; -} - -// Run win_part only (no roundtrip) — compare intermediate output. -static std::vector run_win_part_only(ggml_backend_t backend, - const std::vector& input, - int C, int W, int H, int WS, - int64_t* out_ne) { - size_t ctx_size = ggml_tensor_overhead() * 16 + ggml_graph_overhead(); - struct ggml_init_params params = {ctx_size, nullptr, true}; - struct ggml_context* ctx = ggml_init(params); - - auto* inp = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, C, W, H); - ggml_set_name(inp, "input"); - ggml_set_input(inp); - - auto* parted = ggml_win_part(ctx, inp, WS); - ggml_set_name(parted, "parted"); - ggml_set_output(parted); - - auto* graph = ggml_new_graph(ctx); - ggml_build_forward_expand(graph, parted); - - auto* galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!ggml_gallocr_reserve(galloc, graph) || !ggml_gallocr_alloc_graph(galloc, graph)) { - fprintf(stderr, "FAIL: gallocr_alloc_graph failed\n"); - exit(1); - } - - ggml_backend_tensor_set(inp, input.data(), 0, input.size() * sizeof(float)); - ggml_backend_graph_compute(backend, graph); - - // Output shape - int64_t total = 1; - for (int d = 0; d < 4; d++) { - out_ne[d] = parted->ne[d]; - total *= parted->ne[d]; - } - - std::vector output(total); - ggml_backend_tensor_get(parted, output.data(), 0, output.size() * sizeof(float)); - - ggml_gallocr_free(galloc); - ggml_free(ctx); - return output; -} - -static bool compare_vectors(const std::vector& a, const std::vector& b, - const char* label, float tol = 1e-6f) { - if (a.size() != b.size()) { - fprintf(stderr, " FAIL [%s]: size mismatch: %zu vs %zu\n", label, a.size(), b.size()); - return false; - } - float max_diff = 0.0f; - int first_fail = -1; - for (size_t i = 0; i < a.size(); i++) { - float diff = std::fabs(a[i] - b[i]); - if (diff > max_diff) max_diff = diff; - if (diff > tol && first_fail < 0) first_fail = (int)i; - } - if (first_fail >= 0) { - fprintf(stderr, " FAIL [%s]: max_diff=%.8f at index %d (cpu=%.8f metal=%.8f)\n", - label, max_diff, first_fail, a[first_fail], b[first_fail]); - return false; - } - fprintf(stderr, " PASS [%s]: max_diff=%.2e (%zu elements)\n", label, max_diff, a.size()); - return true; -} - -int main() { -#ifndef GGML_USE_METAL - fprintf(stderr, "Metal not available, skipping test\n"); - return 0; -#else - // Init backends - ggml_backend_t cpu_backend = ggml_backend_cpu_init(); - ggml_backend_t metal_backend = ggml_backend_metal_init(); - if (!metal_backend) { - fprintf(stderr, "FAIL: could not init Metal backend\n"); - return 1; - } - - bool all_pass = true; - - // ── Test 1: sam3 exact dimensions (C=1024, W=H=72, WS=24) ────────── - { - fprintf(stderr, "\n=== Test 1: sam3 dimensions (1024, 72, 72, WS=24) ===\n"); - const int C = 1024, W = 72, H = 72, WS = 24; - std::vector input(C * W * H); - for (size_t i = 0; i < input.size(); i++) { - input[i] = 0.001f * (float)(i % 9973); // pseudo-random - } - - // win_part only - int64_t cpu_ne[4], metal_ne[4]; - auto cpu_part = run_win_part_only(cpu_backend, input, C, W, H, WS, cpu_ne); - auto metal_part = run_win_part_only(metal_backend, input, C, W, H, WS, metal_ne); - fprintf(stderr, " win_part output shape: [%lld, %lld, %lld, %lld]\n", - cpu_ne[0], cpu_ne[1], cpu_ne[2], cpu_ne[3]); - all_pass &= compare_vectors(cpu_part, metal_part, "win_part"); - - // roundtrip (win_part -> win_unpart) - auto cpu_rt = run_win_roundtrip(cpu_backend, input, C, W, H, WS); - auto metal_rt = run_win_roundtrip(metal_backend, input, C, W, H, WS); - all_pass &= compare_vectors(cpu_rt, metal_rt, "roundtrip"); - - // roundtrip should recover original input exactly (no padding needed) - all_pass &= compare_vectors(input, cpu_rt, "cpu_identity"); - all_pass &= compare_vectors(input, metal_rt, "metal_identity"); - } - - // ── Test 2: non-divisible dimensions (needs padding) ──────────────── - { - fprintf(stderr, "\n=== Test 2: padding case (64, 50, 50, WS=24) ===\n"); - const int C = 64, W = 50, H = 50, WS = 24; - std::vector input(C * W * H); - for (size_t i = 0; i < input.size(); i++) { - input[i] = -1.0f + 2.0f * (float)(i % 7919) / 7919.0f; - } - - int64_t cpu_ne[4], metal_ne[4]; - auto cpu_part = run_win_part_only(cpu_backend, input, C, W, H, WS, cpu_ne); - auto metal_part = run_win_part_only(metal_backend, input, C, W, H, WS, metal_ne); - fprintf(stderr, " win_part output shape: [%lld, %lld, %lld, %lld]\n", - cpu_ne[0], cpu_ne[1], cpu_ne[2], cpu_ne[3]); - all_pass &= compare_vectors(cpu_part, metal_part, "win_part_padded"); - - auto cpu_rt = run_win_roundtrip(cpu_backend, input, C, W, H, WS); - auto metal_rt = run_win_roundtrip(metal_backend, input, C, W, H, WS); - all_pass &= compare_vectors(cpu_rt, metal_rt, "roundtrip_padded"); - all_pass &= compare_vectors(input, cpu_rt, "cpu_identity_padded"); - all_pass &= compare_vectors(input, metal_rt, "metal_identity_padded"); - } - - // ── Test 3: small case for easy debugging ────────────────────────── - { - fprintf(stderr, "\n=== Test 3: small (4, 6, 6, WS=3) ===\n"); - const int C = 4, W = 6, H = 6, WS = 3; - std::vector input(C * W * H); - for (size_t i = 0; i < input.size(); i++) { - input[i] = (float)(i + 1); - } - - int64_t cpu_ne[4], metal_ne[4]; - auto cpu_part = run_win_part_only(cpu_backend, input, C, W, H, WS, cpu_ne); - auto metal_part = run_win_part_only(metal_backend, input, C, W, H, WS, metal_ne); - fprintf(stderr, " win_part output shape: [%lld, %lld, %lld, %lld]\n", - cpu_ne[0], cpu_ne[1], cpu_ne[2], cpu_ne[3]); - all_pass &= compare_vectors(cpu_part, metal_part, "win_part_small"); - - auto cpu_rt = run_win_roundtrip(cpu_backend, input, C, W, H, WS); - auto metal_rt = run_win_roundtrip(metal_backend, input, C, W, H, WS); - all_pass &= compare_vectors(cpu_rt, metal_rt, "roundtrip_small"); - all_pass &= compare_vectors(input, cpu_rt, "cpu_identity_small"); - all_pass &= compare_vectors(input, metal_rt, "metal_identity_small"); - } - - // ── Test 4: asymmetric dimensions ────────────────────────────────── - { - fprintf(stderr, "\n=== Test 4: asymmetric (128, 48, 72, WS=24) ===\n"); - const int C = 128, W = 48, H = 72, WS = 24; - std::vector input(C * W * H); - for (size_t i = 0; i < input.size(); i++) { - input[i] = sinf((float)i * 0.01f); - } - - int64_t cpu_ne[4], metal_ne[4]; - auto cpu_part = run_win_part_only(cpu_backend, input, C, W, H, WS, cpu_ne); - auto metal_part = run_win_part_only(metal_backend, input, C, W, H, WS, metal_ne); - fprintf(stderr, " win_part output shape: [%lld, %lld, %lld, %lld]\n", - cpu_ne[0], cpu_ne[1], cpu_ne[2], cpu_ne[3]); - all_pass &= compare_vectors(cpu_part, metal_part, "win_part_asym"); - - auto cpu_rt = run_win_roundtrip(cpu_backend, input, C, W, H, WS); - auto metal_rt = run_win_roundtrip(metal_backend, input, C, W, H, WS); - all_pass &= compare_vectors(cpu_rt, metal_rt, "roundtrip_asym"); - all_pass &= compare_vectors(input, cpu_rt, "cpu_identity_asym"); - all_pass &= compare_vectors(input, metal_rt, "metal_identity_asym"); - } - - ggml_backend_free(cpu_backend); - ggml_backend_free(metal_backend); - - fprintf(stderr, "\n%s\n", all_pass ? "ALL TESTS PASSED" : "SOME TESTS FAILED"); - return all_pass ? 0 : 1; -#endif -} diff --git a/tests/test_phase3.cpp b/tests/test_phase3.cpp deleted file mode 100644 index 3472efe..0000000 --- a/tests/test_phase3.cpp +++ /dev/null @@ -1,632 +0,0 @@ -// Phase 3 Numerical Audit — test ViT backbone, neck, sinusoidal PE, RoPE -// against Python reference tensors dumped by dump_phase3_reference.py. -// -// Usage: -// ./test_phase3 tests/ref_phase3 models/sam3-f16.ggml tests/test_random.jpg -// -// Tests: -// 1. RoPE computation (pure math, no model needed) -// 2. Sinusoidal PE computation (pure math, no model needed) -// 3. Full image encoding: preprocessing → ViT → neck → PE -// Compares intermediate and final outputs against Python references. - -#include "sam3.h" - -#include -#include -#include -#include -#include -#include -#include -#include - -// ═══════════════════════════════════════════════════════════════════════════════ -// Helpers -// ═══════════════════════════════════════════════════════════════════════════════ - -struct ref_tensor { - std::vector data; - std::vector shape; - int numel() const { - int n = 1; - for (auto d : shape) n *= d; - return n; - } -}; - -static ref_tensor load_ref(const std::string & path) { - ref_tensor t; - { - std::ifstream f(path + ".shape"); - if (!f) { - fprintf(stderr, " [SKIP] %s.shape not found\n", path.c_str()); - return t; - } - std::string line; - std::getline(f, line); - size_t pos = 0; - while (pos < line.size()) { - size_t end = line.find(',', pos); - if (end == std::string::npos) end = line.size(); - t.shape.push_back(std::stoi(line.substr(pos, end - pos))); - pos = end + 1; - } - } - { - std::ifstream f(path + ".bin", std::ios::binary); - if (!f) { - fprintf(stderr, " [SKIP] %s.bin not found\n", path.c_str()); - t.shape.clear(); - return t; - } - t.data.resize(t.numel()); - f.read(reinterpret_cast(t.data.data()), t.numel() * sizeof(float)); - } - return t; -} - -struct compare_result { - float max_diff = 0.0f; - float mean_diff = 0.0f; - double cosine_sim = 0.0; - int n_elements = 0; - int n_bad = 0; // elements exceeding tolerance -}; - -static compare_result compare_tensors(const float * a, const float * b, int n, - float atol = 1e-4f) { - compare_result r; - r.n_elements = n; - double sum_diff = 0.0; - double dot_ab = 0.0, dot_aa = 0.0, dot_bb = 0.0; - for (int i = 0; i < n; ++i) { - float diff = fabsf(a[i] - b[i]); - if (diff > r.max_diff) r.max_diff = diff; - if (diff > atol) r.n_bad++; - sum_diff += diff; - dot_ab += (double)a[i] * b[i]; - dot_aa += (double)a[i] * a[i]; - dot_bb += (double)b[i] * b[i]; - } - r.mean_diff = (float)(sum_diff / n); - if (dot_aa > 0 && dot_bb > 0) { - r.cosine_sim = dot_ab / (sqrt(dot_aa) * sqrt(dot_bb)); - } - return r; -} - -static bool check(const std::string & name, const float * got, const ref_tensor & ref, - float atol, int & n_pass, int & n_fail) { - if (ref.data.empty()) { - fprintf(stderr, " [SKIP] %s — no reference data\n", name.c_str()); - return true; - } - auto r = compare_tensors(got, ref.data.data(), ref.numel(), atol); - bool ok = r.max_diff <= atol; - fprintf(stderr, " %s %-45s: max=%.6e mean=%.6e cos=%.8f bad=%d/%d (atol=%.1e)\n", - ok ? "[PASS]" : "[FAIL]", name.c_str(), r.max_diff, r.mean_diff, - r.cosine_sim, r.n_bad, r.n_elements, atol); - if (ok) n_pass++; else n_fail++; - return ok; -} - -// ═══════════════════════════════════════════════════════════════════════════════ -// Test 1: RoPE computation -// ═══════════════════════════════════════════════════════════════════════════════ - -static void test_rope(const std::string & ref_dir, int & n_pass, int & n_fail) { - fprintf(stderr, "\n╔══════════════════════════════════════════╗\n"); - fprintf(stderr, "║ Test: RoPE Frequencies ║\n"); - fprintf(stderr, "╚══════════════════════════════════════════╝\n"); - - struct { int end; float scale; const char * name; } cases[] = { - {24, 1.0f, "rope_window_real"}, - {72, 24.0f/72.0f, "rope_global_real"}, - }; - - const int dim = 64; - const float theta = 10000.0f; - const int half_dim = dim / 4; // 16 - - std::vector freqs(half_dim); - for (int i = 0; i < half_dim; ++i) { - freqs[i] = 1.0f / powf(theta, (float)(i * 4) / (float)dim); - } - - for (auto & tc : cases) { - auto ref = load_ref(ref_dir + "/" + tc.name); - if (ref.data.empty()) continue; - - const int N = tc.end * tc.end; - std::vector our_rope(N * dim); - - for (int idx = 0; idx < N; ++idx) { - float t_x = (float)(idx % tc.end) * tc.scale; - float t_y = static_cast(idx / tc.end) * tc.scale; // NOLINT(bugprone-integer-division) - - for (int i = 0; i < half_dim; ++i) { - float angle_x = t_x * freqs[i]; - our_rope[idx * dim + i * 2 + 0] = cosf(angle_x); - our_rope[idx * dim + i * 2 + 1] = sinf(angle_x); - } - for (int i = 0; i < half_dim; ++i) { - float angle_y = t_y * freqs[i]; - our_rope[idx * dim + half_dim * 2 + i * 2 + 0] = cosf(angle_y); - our_rope[idx * dim + half_dim * 2 + i * 2 + 1] = sinf(angle_y); - } - } - - check(tc.name, our_rope.data(), ref, 1e-5f, n_pass, n_fail); - } - - // Also compare against checkpoint freqs_cis to verify conversion is lossless - for (int bi : {0, 7}) { - std::string name = "ckpt_freqs_cis_block" + std::to_string(bi); - auto ref_ckpt = load_ref(ref_dir + "/" + name); - auto ref_computed = load_ref(ref_dir + "/" + (bi == 0 ? "rope_window_real" : "rope_global_real")); - if (!ref_ckpt.data.empty() && !ref_computed.data.empty()) { - auto r = compare_tensors(ref_ckpt.data.data(), ref_computed.data.data(), - std::min(ref_ckpt.numel(), ref_computed.numel()), 1e-6f); - bool ok = r.max_diff <= 1e-6f; - fprintf(stderr, " %s ckpt vs computed block %d: max=%.6e\n", - ok ? "[PASS]" : "[FAIL]", bi, r.max_diff); - if (ok) n_pass++; else n_fail++; - } - } -} - -// ═══════════════════════════════════════════════════════════════════════════════ -// Test 2: Sinusoidal PE -// ═══════════════════════════════════════════════════════════════════════════════ - -static void test_sinusoidal_pe(const std::string & ref_dir, int & n_pass, int & n_fail) { - fprintf(stderr, "\n╔══════════════════════════════════════════╗\n"); - fprintf(stderr, "║ Test: Sinusoidal PE ║\n"); - fprintf(stderr, "╚══════════════════════════════════════════╝\n"); - - struct { int H; int W; const char * name; } cases[] = { - {288, 288, "pe_288"}, - {144, 144, "pe_144"}, - { 72, 72, "pe_72"}, - { 36, 36, "pe_36"}, - }; - - for (auto & tc : cases) { - auto ref = load_ref(ref_dir + "/" + tc.name); - if (ref.data.empty()) continue; - - const int H = tc.H, W = tc.W, d_model = 256; - const int half = d_model / 2; // 128 - const float scale = 2.0f * (float)M_PI; - const float temperature = 10000.0f; - - // Match Python PositionEmbeddingSine.forward() exactly: - // y_embed = arange(1, H+1) / (H + eps) * scale - // dim_t = temperature ** (2 * (arange(half) // 2) / half) - // pos_y = y_embed / dim_t - // stack(sin(even), cos(odd)).flatten → interleaved sin/cos - // Output: [1, 256, H, W] (NCHW) - - std::vector our_pe(d_model * H * W); - const float eps = 1e-6f; - - for (int y = 0; y < H; ++y) { - for (int x = 0; x < W; ++x) { - float pos_y = ((float)(y + 1) / ((float)H + eps)) * scale; - float pos_x = ((float)(x + 1) / ((float)W + eps)) * scale; - - for (int i = 0; i < half; ++i) { - int paired = (i / 2) * 2; // 0,0,2,2,4,4,... - float dim_t = powf(temperature, (float)paired / (float)half); - - float val_x, val_y; - if (i % 2 == 0) { - val_x = sinf(pos_x / dim_t); - val_y = sinf(pos_y / dim_t); - } else { - val_x = cosf(pos_x / dim_t); - val_y = cosf(pos_y / dim_t); - } - - // PyTorch output is [1, 256, H, W] - // Channel layout: first 128 = pos_y, next 128 = pos_x - our_pe[i * H * W + y * W + x] = val_y; - our_pe[(i + half) * H * W + y * W + x] = val_x; - } - } - } - - check(tc.name, our_pe.data(), ref, 1e-5f, n_pass, n_fail); - } -} - -// ═══════════════════════════════════════════════════════════════════════════════ -// Test 3: Full image encoding (model-dependent) -// ═══════════════════════════════════════════════════════════════════════════════ - -static void test_encode_image(const std::string & model_path, - const std::string & image_path, - const std::string & ref_dir, - int & n_pass, int & n_fail) { - fprintf(stderr, "\n╔══════════════════════════════════════════╗\n"); - fprintf(stderr, "║ Test: Full Image Encoding ║\n"); - fprintf(stderr, "╚══════════════════════════════════════════╝\n"); - - sam3_params params; - params.model_path = model_path; - params.use_gpu = false; - params.n_threads = 4; - - auto model = sam3_load_model(params); - if (!model) { - fprintf(stderr, " FATAL: Failed to load model\n"); - n_fail++; - return; - } - - auto state = sam3_create_state(*model, params); - if (!state) { - fprintf(stderr, " FATAL: Failed to create state\n"); - n_fail++; - return; - } - - auto img = sam3_load_image(image_path); - if (img.data.empty()) { - fprintf(stderr, " FATAL: Failed to load image\n"); - n_fail++; - return; - } - - bool ok = sam3_encode_image(*state, *model, img); - if (!ok) { - fprintf(stderr, " FATAL: sam3_encode_image failed\n"); - n_fail++; - return; - } - - fprintf(stderr, " Image encoding completed successfully\n"); - - std::string dump_dir = ref_dir + "/cpp_out_phase3"; - { - std::string cmd = "mkdir -p " + dump_dir; - (void)system(cmd.c_str()); - } - - // ── Helper: dump and compare a state tensor ────────────────────────── - // C++ tensors are in ggml layout. We need to transpose to PyTorch layout. - - // For NHWC tensors (ViT intermediates): ggml [E, W, H, B] → PyTorch [B, H, W, E] - auto compare_nhwc = [&](const std::string & cpp_name, const std::string & ref_name, - float atol) { - bool dumped = sam3_dump_state_tensor(*state, cpp_name, dump_dir + "/" + cpp_name); - if (!dumped) { - fprintf(stderr, " [SKIP] %s — tensor '%s' not available in state\n", - ref_name.c_str(), cpp_name.c_str()); - return; - } - auto ref_t = load_ref(ref_dir + "/" + ref_name); - auto cpp_t = load_ref(dump_dir + "/" + cpp_name); - if (ref_t.data.empty() || cpp_t.data.empty()) { - fprintf(stderr, " [SKIP] %s — data not available\n", ref_name.c_str()); - return; - } - // ggml [E, W, H, ...] → [H, W, E] (ignoring batch) - int E = cpp_t.shape[0]; - int nW = cpp_t.shape[1]; - int nH = cpp_t.shape.size() > 2 ? cpp_t.shape[2] : 1; - // ref is [1, H, W, E] or [B, H, W, E] - std::vector transposed(cpp_t.numel()); - for (int h = 0; h < nH; ++h) - for (int w = 0; w < nW; ++w) - for (int e = 0; e < E; ++e) - transposed[h * nW * E + w * E + e] = cpp_t.data[e + w * E + h * E * nW]; - - check(ref_name + " (NHWC)", transposed.data(), ref_t, atol, n_pass, n_fail); - }; - - // For NCHW tensors (neck outputs): ggml [C, W, H, B] → PyTorch [1, C, H, W] - auto compare_nchw = [&](const std::string & cpp_name, const std::string & ref_name, - float atol) { - bool dumped = sam3_dump_state_tensor(*state, cpp_name, dump_dir + "/" + cpp_name); - if (!dumped) { - fprintf(stderr, " [SKIP] %s — tensor '%s' not available in state\n", - ref_name.c_str(), cpp_name.c_str()); - return; - } - auto ref_t = load_ref(ref_dir + "/" + ref_name); - auto cpp_t = load_ref(dump_dir + "/" + cpp_name); - if (ref_t.data.empty() || cpp_t.data.empty()) { - fprintf(stderr, " [SKIP] %s — data not available\n", ref_name.c_str()); - return; - } - // ggml [C, W, H, ...]: element (c,w,h) at c + w*C + h*C*W - // PyTorch [1, C, H, W]: element (0,c,h,w) at c*H*W + h*W + w - int C = cpp_t.shape[0]; - int nW = cpp_t.shape[1]; - int nH = cpp_t.shape.size() > 2 ? cpp_t.shape[2] : 1; - std::vector transposed(cpp_t.numel()); - for (int c = 0; c < C; ++c) - for (int h = 0; h < nH; ++h) - for (int w = 0; w < nW; ++w) - transposed[c * nH * nW + h * nW + w] = cpp_t.data[c + w * C + h * C * nW]; - - check(ref_name + " (NCHW)", transposed.data(), ref_t, atol, n_pass, n_fail); - }; - - // ── Compare ViT intermediates ──────────────────────────────────────── - fprintf(stderr, "\n --- ViT Intermediates ---\n"); - - // Patch embed: max_diff ~0.3 due to different bilinear resize between C++ and Python. - // This is the root cause of all downstream differences. - compare_nhwc("dbg_patch_embed", "patch_embed", 0.5f); - - // After pos embed: same as patch_embed (pos embed loaded exactly) - compare_nhwc("dbg_after_pos_embed", "after_pos_embed", 0.5f); - - // After ln_pre: amplified by layer norm (max ~1.2) - compare_nhwc("dbg_after_ln_pre", "after_ln_pre", 2.0f); - - // Block outputs (error accumulates through transformer layers) - compare_nhwc("dbg_block_0_out", "block_0_out", 2.0f); - compare_nhwc("dbg_block_7_out", "block_7_out", 5.0f); - - // ViT final output - sam3_dump_state_tensor(*state, "vit_output", dump_dir + "/vit_output"); - auto ref_vit = load_ref(ref_dir + "/vit_output_bchw"); - auto cpp_vit = load_ref(dump_dir + "/vit_output"); - if (!ref_vit.data.empty() && !cpp_vit.data.empty()) { - int E = cpp_vit.shape[0], W = cpp_vit.shape[1]; - int H = (cpp_vit.shape.size() > 2) ? cpp_vit.shape[2] : 1; - std::vector transposed(cpp_vit.numel()); - for (int e = 0; e < E; ++e) - for (int h = 0; h < H; ++h) - for (int w = 0; w < W; ++w) - transposed[e * H * W + h * W + w] = cpp_vit.data[e + w * E + h * E * W]; - // ViT output after 32 blocks: max_diff can be ~80 due to preprocessing diff - // (C++ bilinear resize vs Python PIL resize) amplified through 32 transformer layers. - // Cosine sim >0.995 confirms the ViT is computing the correct function. - // The isolated test (Test 4) with identical input achieves cos=1.0 on f32. - auto vit_r = compare_tensors(transposed.data(), ref_vit.data.data(), ref_vit.numel(), 100.0f); - bool vit_ok = vit_r.cosine_sim > 0.995 && vit_r.mean_diff < 0.1f; - fprintf(stderr, " %s %-45s: max=%.6e mean=%.6e cos=%.8f (tol: cos>0.995, mean<0.1)\n", - vit_ok ? "[PASS]" : "[FAIL]", "vit_output_bchw", - vit_r.max_diff, vit_r.mean_diff, vit_r.cosine_sim); - if (vit_ok) n_pass++; else n_fail++; - } - - // ── Compare neck outputs ───────────────────────────────────────────── - fprintf(stderr, "\n --- Neck Outputs (Detector) ---\n"); - for (int i = 0; i < 4; ++i) { - std::string name = "neck_det_" + std::to_string(i); - compare_nchw(name, name, 5.0f); - } - - // ── Compare sinusoidal PE from state ───────────────────────────────── - // The state stores PE tensors as neck_det_pe_i. - fprintf(stderr, "\n --- Sinusoidal PE (from state) ---\n"); - int pe_sizes[] = {288, 144, 72, 36}; - for (int i = 0; i < 4; ++i) { - std::string cpp_pe_name = "neck_det_pe_" + std::to_string(i); - std::string ref_pe_name = "pe_" + std::to_string(pe_sizes[i]); - sam3_dump_state_tensor(*state, cpp_pe_name, dump_dir + "/" + cpp_pe_name); - auto ref_pe = load_ref(ref_dir + "/" + ref_pe_name); - auto cpp_pe = load_ref(dump_dir + "/" + cpp_pe_name); - if (!ref_pe.data.empty() && !cpp_pe.data.empty()) { - // PE is stored as [D, W, H, 1] in ggml, ref is [1, 256, H, W] in PyTorch - int D = cpp_pe.shape[0], pW = cpp_pe.shape[1]; - int pH = (cpp_pe.shape.size() > 2) ? cpp_pe.shape[2] : 1; - std::vector transposed(cpp_pe.numel()); - for (int d = 0; d < D; ++d) - for (int h = 0; h < pH; ++h) - for (int w = 0; w < pW; ++w) - transposed[d * pH * pW + h * pW + w] = cpp_pe.data[d + w * D + h * D * pW]; - check(ref_pe_name + " (state)", transposed.data(), ref_pe, 1e-4f, n_pass, n_fail); - } - } - - state.reset(); - sam3_free_model(*model); - model.reset(); -} - -// ═══════════════════════════════════════════════════════════════════════════════ -// Test 4: Isolated ViT numerics using Python-preprocessed input -// This feeds the exact same preprocessed image to the C++ pipeline to isolate -// ViT numerical accuracy from preprocessing differences. -// ═══════════════════════════════════════════════════════════════════════════════ - -static void test_encode_from_preprocessed(const std::string & model_path, - const std::string & ref_dir, - int & n_pass, int & n_fail) { - fprintf(stderr, "\n╔══════════════════════════════════════════╗\n"); - fprintf(stderr, "║ Test: ViT Numerics (Python preprocess) ║\n"); - fprintf(stderr, "╚══════════════════════════════════════════╝\n"); - - // Load Python-preprocessed image (CHW layout, [1, 3, 1008, 1008]) - auto preproc = load_ref(ref_dir + "/preprocessed"); - if (preproc.data.empty()) { - fprintf(stderr, " [SKIP] No preprocessed reference data\n"); - return; - } - fprintf(stderr, " Loaded preprocessed image: shape=["); - for (size_t i = 0; i < preproc.shape.size(); ++i) { - if (i > 0) fprintf(stderr, ", "); - fprintf(stderr, "%d", preproc.shape[i]); - } - fprintf(stderr, "] numel=%d\n", preproc.numel()); - - // The preprocessed data is [1, 3, 1008, 1008] in PyTorch NCHW layout. - // For ggml, the input tensor is [W=1008, H=1008, C=3, B=1]. - // ggml flat index: x + y*W + c*W*H = same as CHW[c*H*W + y*W + x]. - // So the CHW data can be uploaded directly. - const float * chw_data = preproc.data.data(); - int img_size = preproc.shape[2]; // 1008 - - sam3_params params; - params.model_path = model_path; - params.use_gpu = false; - params.n_threads = 4; - - auto model = sam3_load_model(params); - if (!model) { - fprintf(stderr, " FATAL: Failed to load model\n"); - n_fail++; - return; - } - - auto state = sam3_create_state(*model, params); - if (!state) { - fprintf(stderr, " FATAL: Failed to create state\n"); - n_fail++; - return; - } - - bool ok = sam3_encode_image_from_preprocessed(*state, *model, chw_data, img_size); - if (!ok) { - fprintf(stderr, " FATAL: sam3_encode_image_from_preprocessed failed\n"); - n_fail++; - return; - } - - std::string dump_dir = ref_dir + "/cpp_out_from_preproc"; - { - std::string cmd = "mkdir -p " + dump_dir; - (void)system(cmd.c_str()); - } - - // ── Compare ViT output ──────────────────────────────────────────────── - fprintf(stderr, "\n --- ViT Output (same input) ---\n"); - sam3_dump_state_tensor(*state, "vit_output", dump_dir + "/vit_output"); - auto ref_vit = load_ref(ref_dir + "/vit_output_bchw"); - auto cpp_vit = load_ref(dump_dir + "/vit_output"); - if (!ref_vit.data.empty() && !cpp_vit.data.empty()) { - int E = cpp_vit.shape[0], W = cpp_vit.shape[1]; - int H = (cpp_vit.shape.size() > 2) ? cpp_vit.shape[2] : 1; - std::vector transposed(cpp_vit.numel()); - for (int e = 0; e < E; ++e) - for (int h = 0; h < H; ++h) - for (int w = 0; w < W; ++w) - transposed[e * H * W + h * W + w] = cpp_vit.data[e + w * E + h * E * W]; - - auto vit_r = compare_tensors(transposed.data(), ref_vit.data.data(), ref_vit.numel(), 1e-2f); - // With identical input and f32 weights, expect very high cosine similarity. - // With f16 weights, expect cos > 0.999 and mean_diff < 0.05. - bool vit_ok = vit_r.cosine_sim > 0.999 && vit_r.mean_diff < 0.05f; - fprintf(stderr, " %s %-45s: max=%.6e mean=%.6e cos=%.8f bad=%d/%d\n", - vit_ok ? "[PASS]" : "[FAIL]", "vit_output (same input)", - vit_r.max_diff, vit_r.mean_diff, vit_r.cosine_sim, - vit_r.n_bad, vit_r.n_elements); - if (vit_ok) n_pass++; else n_fail++; - } - - // ── Compare neck outputs ───────────────────────────────────────────── - fprintf(stderr, "\n --- Neck Outputs (same input) ---\n"); - const char * neck_names[] = {"neck_det_0", "neck_det_1", "neck_det_2", "neck_det_3"}; - for (int i = 0; i < 4; ++i) { - bool dumped = sam3_dump_state_tensor(*state, neck_names[i], dump_dir + "/" + neck_names[i]); - if (!dumped) { - fprintf(stderr, " [SKIP] %s — not available\n", neck_names[i]); - continue; - } - auto ref_t = load_ref(ref_dir + "/" + neck_names[i]); - auto cpp_t = load_ref(dump_dir + "/" + neck_names[i]); - if (ref_t.data.empty() || cpp_t.data.empty()) continue; - - int C = cpp_t.shape[0]; - int nW = cpp_t.shape[1]; - int nH = cpp_t.shape.size() > 2 ? cpp_t.shape[2] : 1; - std::vector transposed(cpp_t.numel()); - for (int c = 0; c < C; ++c) - for (int h = 0; h < nH; ++h) - for (int w = 0; w < nW; ++w) - transposed[c * nH * nW + h * nW + w] = cpp_t.data[c + w * C + h * C * nW]; - - auto r = compare_tensors(transposed.data(), ref_t.data.data(), ref_t.numel(), 1.0f); - bool pass = r.cosine_sim > 0.999 && r.mean_diff < 0.05f; - fprintf(stderr, " %s %-45s: max=%.6e mean=%.6e cos=%.8f bad=%d/%d\n", - pass ? "[PASS]" : "[FAIL]", neck_names[i], - r.max_diff, r.mean_diff, r.cosine_sim, r.n_bad, r.n_elements); - if (pass) n_pass++; else n_fail++; - } - - // ── Compare tracker neck outputs ──────────────────────────────────── - fprintf(stderr, "\n --- Tracker Neck Outputs (same input) ---\n"); - const char * trk_names[] = {"neck_trk_0", "neck_trk_1", "neck_trk_2", "neck_trk_3"}; - for (int i = 0; i < 4; ++i) { - bool dumped = sam3_dump_state_tensor(*state, trk_names[i], dump_dir + "/" + trk_names[i]); - if (!dumped) { - fprintf(stderr, " [SKIP] %s — not available\n", trk_names[i]); - continue; - } - auto ref_t = load_ref(ref_dir + "/" + trk_names[i]); - auto cpp_t = load_ref(dump_dir + "/" + trk_names[i]); - if (ref_t.data.empty() || cpp_t.data.empty()) continue; - - int C = cpp_t.shape[0]; - int nW = cpp_t.shape[1]; - int nH = cpp_t.shape.size() > 2 ? cpp_t.shape[2] : 1; - std::vector transposed(cpp_t.numel()); - for (int c = 0; c < C; ++c) - for (int h = 0; h < nH; ++h) - for (int w = 0; w < nW; ++w) - transposed[c * nH * nW + h * nW + w] = cpp_t.data[c + w * C + h * C * nW]; - - auto r = compare_tensors(transposed.data(), ref_t.data.data(), ref_t.numel(), 1.0f); - bool pass = r.cosine_sim > 0.999 && r.mean_diff < 0.05f; - fprintf(stderr, " %s %-45s: max=%.6e mean=%.6e cos=%.8f bad=%d/%d\n", - pass ? "[PASS]" : "[FAIL]", trk_names[i], - r.max_diff, r.mean_diff, r.cosine_sim, r.n_bad, r.n_elements); - if (pass) n_pass++; else n_fail++; - } - - state.reset(); - sam3_free_model(*model); - model.reset(); -} - -// ═══════════════════════════════════════════════════════════════════════════════ -// Main -// ═══════════════════════════════════════════════════════════════════════════════ - -int main(int argc, char ** argv) { - if (argc < 2) { - fprintf(stderr, "Usage: %s [model.ggml] [image.jpg]\n", argv[0]); - fprintf(stderr, "\nGenerate reference data first:\n"); - fprintf(stderr, " uv run python tests/dump_phase3_reference.py \\\n"); - fprintf(stderr, " --checkpoint raw_weights/sam3.pt --image data/test_image.jpg\n"); - return 1; - } - - std::string ref_dir = argv[1]; - std::string model_path = (argc > 2) ? argv[2] : ""; - std::string image_path = (argc > 3) ? argv[3] : ""; - - int n_pass = 0, n_fail = 0; - - // Test 1: RoPE (no model needed) - test_rope(ref_dir, n_pass, n_fail); - - // Test 2: Sinusoidal PE (no model needed) - test_sinusoidal_pe(ref_dir, n_pass, n_fail); - - // Test 3: Full encoding (needs model + image) - if (!model_path.empty() && !image_path.empty()) { - test_encode_image(model_path, image_path, ref_dir, n_pass, n_fail); - } else { - fprintf(stderr, "\n[SKIP] Full encoding test — no model/image provided\n"); - } - - // Test 4: Isolated ViT numerics from Python-preprocessed input (needs model) - if (!model_path.empty()) { - test_encode_from_preprocessed(model_path, ref_dir, n_pass, n_fail); - } else { - fprintf(stderr, "\n[SKIP] Isolated ViT test — no model provided\n"); - } - - fprintf(stderr, "\n═══════════════════════════════════════════\n"); - fprintf(stderr, "Phase 3 Audit Results: %d passed, %d failed\n", n_pass, n_fail); - fprintf(stderr, "═══════════════════════════════════════════\n"); - return n_fail > 0 ? 1 : 0; -} diff --git a/tests/test_phase4.cpp b/tests/test_phase4.cpp deleted file mode 100644 index 1ad6424..0000000 --- a/tests/test_phase4.cpp +++ /dev/null @@ -1,334 +0,0 @@ -#include "sam3.h" - -#include "test_utils.h" - -#include -#include -#include -#include -#include - -struct prompt_case { - std::string id; - std::string text; -}; - -struct summary_row { - bool measured = false; - bool ok = true; - float max_diff = 0.0f; - float mean_diff = 0.0f; - float cosine_sim = 1.0f; - std::string worst_tensor; - std::string note; -}; - -static bool ends_with(const std::string & s, const std::string & suffix) { - return s.size() >= suffix.size() && - s.compare(s.size() - suffix.size(), suffix.size(), suffix) == 0; -} - -static std::vector load_prompt_cases(const std::string & path) { - std::vector cases; - std::ifstream f(path); - std::string line; - while (std::getline(f, line)) { - if (line.empty()) { - continue; - } - - prompt_case pc; - size_t tab = line.find('\t'); - if (tab == std::string::npos) { - pc.id = "prompt_" + std::to_string(cases.size()); - pc.text = line; - } else { - pc.id = line.substr(0, tab); - pc.text = line.substr(tab + 1); - } - cases.push_back(std::move(pc)); - } - return cases; -} - -static std::string category_for_tensor(const std::string & name) { - if (name == "causal_mask") return "causal_mask"; - if (name == "text_token_embed") return "text_token_embed"; - if (name == "text_after_pos_embed") return "text_after_pos_embed"; - if (name == "text_final_ln") return "text_final_ln"; - if (name == "text_features_2d") return "text_features_2d"; - - if (ends_with(name, "_after_ln1")) return "block_after_ln1"; - if (ends_with(name, "_qkv")) return "block_qkv"; - if (ends_with(name, "_attn_out")) return "block_attn_out"; - if (ends_with(name, "_after_attn_residual")) return "block_after_attn_residual"; - if (ends_with(name, "_after_ln2")) return "block_after_ln2"; - if (ends_with(name, "_mlp_fc1")) return "block_mlp_fc1"; - if (ends_with(name, "_mlp_gelu")) return "block_mlp_gelu"; - if (ends_with(name, "_mlp_out")) return "block_mlp_out"; - if (ends_with(name, "_out")) return "block_out"; - - return name; -} - -static const char * label_for_category(const std::string & category) { - if (category == "token_ids") return "Token IDs"; - if (category == "causal_mask") return "Causal mask"; - if (category == "text_token_embed") return "Token embedding"; - if (category == "text_after_pos_embed") return "Embedding + position"; - if (category == "block_after_ln1") return "Blocks after ln1"; - if (category == "block_qkv") return "Blocks qkv"; - if (category == "block_attn_out") return "Blocks attn out"; - if (category == "block_after_attn_residual") return "Blocks attn residual"; - if (category == "block_after_ln2") return "Blocks after ln2"; - if (category == "block_mlp_fc1") return "Blocks mlp fc1"; - if (category == "block_mlp_gelu") return "Blocks mlp gelu"; - if (category == "block_mlp_out") return "Blocks mlp out"; - if (category == "block_out") return "Blocks output"; - if (category == "text_final_ln") return "Final layer norm"; - if (category == "text_features_2d") return "Final text features"; - return category.c_str(); -} - -static std::vector tensor_names() { - std::vector names = { - "causal_mask", - "text_token_embed", - "text_after_pos_embed", - "text_final_ln", - "text_features_2d", - }; - - for (int i = 0; i < 24; ++i) { - char name[64]; - - snprintf(name, sizeof(name), "text_block_%02d_after_ln1", i); - names.emplace_back(name); - - snprintf(name, sizeof(name), "text_block_%02d_qkv", i); - names.emplace_back(name); - - snprintf(name, sizeof(name), "text_block_%02d_attn_out", i); - names.emplace_back(name); - - snprintf(name, sizeof(name), "text_block_%02d_after_attn_residual", i); - names.emplace_back(name); - - snprintf(name, sizeof(name), "text_block_%02d_after_ln2", i); - names.emplace_back(name); - - snprintf(name, sizeof(name), "text_block_%02d_mlp_fc1", i); - names.emplace_back(name); - - snprintf(name, sizeof(name), "text_block_%02d_mlp_gelu", i); - names.emplace_back(name); - - snprintf(name, sizeof(name), "text_block_%02d_mlp_out", i); - names.emplace_back(name); - - snprintf(name, sizeof(name), "text_block_%02d_out", i); - names.emplace_back(name); - } - - return names; -} - -static void update_summary(summary_row & row, - const std::string & tensor_name, - const compare_result & r, - bool ok) { - if (!row.measured || r.max_diff >= row.max_diff) { - row.max_diff = r.max_diff; - row.mean_diff = r.mean_diff; - row.cosine_sim = r.cosine_sim; - row.worst_tensor = tensor_name; - } - row.measured = true; - row.ok = row.ok && ok; -} - -int main(int argc, char ** argv) { - if (argc < 4) { - fprintf(stderr, "Usage: %s \n", argv[0]); - return 1; - } - - const std::string ref_dir = argv[1]; - const std::string model_path = argv[2]; - const std::string prompts_path = argv[3]; - const float atol = 1e-4f; - - auto prompts = load_prompt_cases(prompts_path); - if (prompts.empty()) { - fprintf(stderr, "No prompts loaded from %s\n", prompts_path.c_str()); - return 1; - } - - if (!sam3_test_load_tokenizer(model_path)) { - fprintf(stderr, "Failed to load tokenizer from %s\n", model_path.c_str()); - return 1; - } - - sam3_params params; - params.model_path = model_path; - params.n_threads = 1; - params.use_gpu = false; - - auto model = sam3_load_model(params); - if (!model) { - fprintf(stderr, "Failed to load model from %s\n", model_path.c_str()); - return 1; - } - - const std::string cpp_root = ref_dir + "/cpp_out_phase4"; - if (!ensure_dir(cpp_root)) { - fprintf(stderr, "Failed to create %s\n", cpp_root.c_str()); - return 1; - } - - const std::vector names = tensor_names(); - const std::vector category_order = { - "token_ids", - "causal_mask", - "text_token_embed", - "text_after_pos_embed", - "block_after_ln1", - "block_qkv", - "block_attn_out", - "block_after_attn_residual", - "block_after_ln2", - "block_mlp_fc1", - "block_mlp_gelu", - "block_mlp_out", - "block_out", - "text_final_ln", - "text_features_2d", - }; - - std::map> report; - bool overall_ok = true; - - for (const auto & prompt : prompts) { - fprintf(stderr, "\n=== Prompt %s: \"%s\" ===\n", prompt.id.c_str(), prompt.text.c_str()); - - const std::string prompt_ref_dir = ref_dir + "/" + prompt.id; - const std::string prompt_cpp_dir = cpp_root + "/" + prompt.id; - if (!ensure_dir(prompt_cpp_dir)) { - fprintf(stderr, "Failed to create %s\n", prompt_cpp_dir.c_str()); - overall_ok = false; - continue; - } - - auto ref_token_ids = load_ref_i32(prompt_ref_dir + "/token_ids"); - if (ref_token_ids.data.empty()) { - fprintf(stderr, "Missing reference token_ids for %s\n", prompt.id.c_str()); - overall_ok = false; - continue; - } - - auto cpp_token_ids = sam3_test_tokenize(prompt.text); - int token_bad = compare_exact_i32(cpp_token_ids, ref_token_ids.data); - summary_row token_row; - token_row.measured = true; - token_row.ok = (token_bad == 0); - token_row.max_diff = token_bad == 0 ? 0.0f : 1.0f; - token_row.mean_diff = token_bad == 0 ? 0.0f : 1.0f; - token_row.cosine_sim = token_bad == 0 ? 1.0f : 0.0f; - token_row.note = token_bad == 0 ? "exact int compare" : "tokenizer mismatch"; - report[prompt.id]["token_ids"] = token_row; - if (token_bad != 0) { - fprintf(stderr, " [FAIL] token_ids mismatch for %s\n", prompt.id.c_str()); - overall_ok = false; - } - - if (!sam3_test_dump_text_encoder(*model, ref_token_ids.data, prompt_cpp_dir, params.n_threads)) { - fprintf(stderr, " [FAIL] could not dump C++ text encoder tensors for %s\n", prompt.id.c_str()); - overall_ok = false; - continue; - } - - for (const auto & name : names) { - auto ref = load_ref_f32(prompt_ref_dir + "/" + name); - auto got = load_ref_f32(prompt_cpp_dir + "/" + name); - const std::string category = category_for_tensor(name); - - if (ref.data.empty() || got.data.empty()) { - auto & row = report[prompt.id][category]; - row.ok = false; - row.note = "missing reference or C++ output"; - overall_ok = false; - fprintf(stderr, " [FAIL] missing tensor %s for %s\n", name.c_str(), prompt.id.c_str()); - continue; - } - - if (ref.numel() != got.numel()) { - auto & row = report[prompt.id][category]; - row.ok = false; - row.note = "shape mismatch"; - overall_ok = false; - fprintf(stderr, " [FAIL] shape mismatch for %s (%d vs %d)\n", - name.c_str(), got.numel(), ref.numel()); - continue; - } - - auto r = compare_tensors(got.data.data(), ref.data.data(), ref.numel(), atol); - bool ok = r.max_diff <= atol; - update_summary(report[prompt.id][category], name, r, ok); - if (!ok) { - overall_ok = false; - fprintf(stderr, - " [FAIL] %-32s max=%.3e mean=%.3e cosine=%.8f\n", - name.c_str(), r.max_diff, r.mean_diff, r.cosine_sim); - } - } - - for (const auto & category : category_order) { - auto it = report[prompt.id].find(category); - if (it == report[prompt.id].end() || !it->second.measured) { - if (category == "token_ids") { - continue; - } - fprintf(stderr, " [FAIL] no measurements recorded for %s/%s\n", - prompt.id.c_str(), category.c_str()); - overall_ok = false; - report[prompt.id][category].ok = false; - report[prompt.id][category].note = "untested"; - } - } - } - - fprintf(stderr, "\n## Numerical Precision Report\n\n"); - fprintf(stderr, "| Prompt | Operation | max_diff | mean_diff | cosine | Tolerance | Status | Notes |\n"); - fprintf(stderr, "|--------|-----------|----------|-----------|--------|-----------|--------|-------|\n"); - for (const auto & prompt : prompts) { - for (const auto & category : category_order) { - const auto it = report[prompt.id].find(category); - summary_row row; - if (it != report[prompt.id].end()) { - row = it->second; - } - - const char * status = !row.measured && category != "token_ids" - ? "UNTESTED" - : (row.ok ? "PASS" : "FAIL"); - - std::string note = row.note; - if (note.empty() && !row.worst_tensor.empty()) { - note = "worst=" + row.worst_tensor; - } - - fprintf(stderr, - "| %s | %s | %.3e | %.3e | %.8f | %.1e | %s | %s |\n", - prompt.id.c_str(), - label_for_category(category), - row.max_diff, - row.mean_diff, - row.cosine_sim, - atol, - status, - note.empty() ? "measured" : note.c_str()); - } - } - - return overall_ok ? 0 : 1; -} diff --git a/tests/test_phase5.cpp b/tests/test_phase5.cpp deleted file mode 100644 index b9a9e3d..0000000 --- a/tests/test_phase5.cpp +++ /dev/null @@ -1,256 +0,0 @@ -#include "sam3.h" - -#include "test_utils.h" - -#include -#include -#include -#include - -struct tensor_case { - std::string name; - std::string label; -}; - -struct report_row { - bool measured = false; - bool ok = false; - compare_result r; - std::string note; -}; - -static std::string load_prompt(const std::string & path) { - std::ifstream f(path); - std::string line; - if (f) { - std::getline(f, line); - } - return line; -} - -static std::vector tensor_cases() { - std::vector cases = { - {"text_valid_mask", "Text valid mask"}, - {"text_features", "Text features"}, - {"fenc_img_input", "Fusion encoder image input"}, - {"fenc_pos_embed", "Fusion encoder position input"}, - {"fenc_prompt", "Fusion encoder prompt input"}, - {"img_pe_72", "Image PE 72x72"}, - {"fenc_output", "Fusion encoder final output"}, - {"ddec_query_embed", "Decoder query embed"}, - {"ddec_ref_pts_raw", "Decoder reference points raw"}, - {"ddec_presence_token", "Decoder presence token"}, - {"ddec_ref_boxes_init", "Decoder initial reference boxes"}, - {"ddec_query_sine_0", "Decoder layer 0 query sine"}, - {"ddec_query_pos_0", "Decoder layer 0 query pos"}, - {"ddec_rpb_mask_0", "Decoder layer 0 RPB mask"}, - {"ddec_layer0_after_sa", "Decoder layer 0 after self-attn"}, - {"ddec_layer0_after_text_ca", "Decoder layer 0 after text CA"}, - {"ddec_layer0_after_img_ca", "Decoder layer 0 after image CA"}, - {"ddec_layer0_full_out", "Decoder layer 0 full output"}, - {"ddec_layer0_presence", "Decoder layer 0 presence"}, - {"ddec_normed_output", "Decoder final normed output"}, - {"ddec_pred_boxes", "Decoder predicted boxes"}, - {"ddec_presence_logit", "Decoder presence logit"}, - {"scoring_prompt_mlp_out", "Scoring prompt MLP"}, - {"scoring_pooled", "Scoring pooled prompt"}, - {"scoring_proj_pooled", "Scoring projected pooled prompt"}, - {"scoring_proj_hs", "Scoring projected queries"}, - {"scoring_class_scores", "Scoring class scores"}, - {"seg_enc_after_ca", "Seg head encoder after prompt CA"}, - {"seg_enc_visual", "Seg head visual encoder map"}, - {"seg_pixel_dec_stage0", "Seg head pixel decoder stage 0"}, - {"seg_pixel_dec_stage1", "Seg head pixel decoder stage 1"}, - {"seg_pixel_decoder_out", "Seg head pixel decoder output"}, - {"seg_instance_embed", "Seg head instance embed"}, - {"seg_mask_embed", "Seg head mask embed"}, - {"seg_mask_logits", "Seg head mask logits"}, - }; - - for (int i = 0; i < 6; ++i) { - char name[64]; - char label[64]; - - snprintf(name, sizeof(name), "fenc_layer%d_out", i); - snprintf(label, sizeof(label), "Fusion encoder layer %d", i); - cases.push_back({name, label}); - - snprintf(name, sizeof(name), "ddec_layer%d_out", i); - snprintf(label, sizeof(label), "Decoder layer %d output", i); - cases.push_back({name, label}); - - snprintf(name, sizeof(name), "ddec_layer%d_refboxes", i); - snprintf(label, sizeof(label), "Decoder layer %d ref boxes", i); - cases.push_back({name, label}); - } - - return cases; -} - -static float tolerance_for(const std::string & name) { - if (name == "ddec_rpb_mask_0") { - return 2e-4f; - } - if (name == "seg_mask_embed") { - return 2e-4f; - } - if (name == "seg_mask_logits") { - return 3e-3f; - } - if (name == "ddec_layer2_out" || - name == "ddec_layer3_out" || - name == "ddec_layer4_out" || - name == "ddec_layer5_out") { - return 2e-4f; - } - return 1e-4f; -} - -int main(int argc, char ** argv) { - if (argc < 4) { - fprintf(stderr, "Usage: %s \n", argv[0]); - return 1; - } - - const std::string ref_dir = argv[1]; - const std::string model_path = argv[2]; - const std::string image_path = argv[3]; - const std::string cpp_dir = ref_dir + "/cpp_out_phase5"; - if (!ensure_dir(cpp_dir)) { - fprintf(stderr, "Failed to create %s\n", cpp_dir.c_str()); - return 1; - } - - std::string prompt = load_prompt(ref_dir + "/prompt.txt"); - if (prompt.empty()) { - prompt = "yellow school bus"; - } - - auto ref_token_ids = load_ref_i32(ref_dir + "/token_ids"); - - if (!sam3_test_load_tokenizer(model_path)) { - fprintf(stderr, "Failed to load tokenizer from %s\n", model_path.c_str()); - return 1; - } - - auto cpp_token_ids = sam3_test_tokenize(prompt); - std::vector run_token_ids = cpp_token_ids; - - report_row token_row; - if (!ref_token_ids.data.empty()) { - const int token_bad = compare_exact_i32(cpp_token_ids, ref_token_ids.data); - token_row.measured = true; - token_row.ok = token_bad == 0; - token_row.r.max_diff = token_bad == 0 ? 0.0f : 1.0f; - token_row.r.mean_diff = token_bad == 0 ? 0.0f : 1.0f; - token_row.r.cosine_sim = token_bad == 0 ? 1.0f : 0.0f; - token_row.note = token_bad == 0 ? "exact int compare" : "tokenizer mismatch"; - run_token_ids = ref_token_ids.data; - } else { - token_row.measured = false; - token_row.ok = false; - token_row.note = "missing token_ids reference"; - } - - sam3_params params; - params.model_path = model_path; - params.n_threads = 1; - params.use_gpu = false; - - auto model = sam3_load_model(params); - if (!model) { - fprintf(stderr, "Failed to load model from %s\n", model_path.c_str()); - return 1; - } - - (void) image_path; - - const std::string prephase_ref_dir = ref_dir + "/../ref"; - if (!sam3_test_dump_phase5_from_ref_inputs(*model, run_token_ids, - prephase_ref_dir, ref_dir, - cpp_dir, params.n_threads)) { - fprintf(stderr, "Failed to dump C++ phase 5 tensors\n"); - return 1; - } - - const auto cases = tensor_cases(); - std::vector rows(cases.size()); - bool overall_ok = token_row.measured ? token_row.ok : false; - - for (size_t i = 0; i < cases.size(); ++i) { - const auto & tc = cases[i]; - auto ref = load_ref_f32(ref_dir + "/" + tc.name); - auto got = load_ref_f32(cpp_dir + "/" + tc.name); - - if (ref.data.empty() || got.data.empty()) { - rows[i].ok = false; - rows[i].note = "missing reference or C++ output"; - overall_ok = false; - fprintf(stderr, " [FAIL] missing tensor %s\n", tc.name.c_str()); - continue; - } - - if (ref.numel() != got.numel()) { - rows[i].ok = false; - rows[i].note = "numel mismatch"; - overall_ok = false; - fprintf(stderr, " [FAIL] numel mismatch for %s (%d vs %d)\n", - tc.name.c_str(), got.numel(), ref.numel()); - continue; - } - - rows[i].measured = true; - const float atol = tolerance_for(tc.name); - rows[i].r = compare_tensors(got.data.data(), ref.data.data(), ref.numel(), atol); - rows[i].ok = rows[i].r.max_diff <= atol; - if (!rows[i].ok) { - char note[160]; - snprintf(note, sizeof(note), - "worst_idx=%d got=%.6g ref=%.6g bad=%d/%d", - rows[i].r.worst_index, - rows[i].r.worst_a, - rows[i].r.worst_b, - rows[i].r.n_bad, - rows[i].r.n_total); - rows[i].note = note; - overall_ok = false; - fprintf(stderr, - " [FAIL] %-24s max=%.3e mean=%.3e cosine=%.8f idx=%d got=%.6g ref=%.6g\n", - tc.name.c_str(), - rows[i].r.max_diff, - rows[i].r.mean_diff, - rows[i].r.cosine_sim, - rows[i].r.worst_index, - rows[i].r.worst_a, - rows[i].r.worst_b); - } - } - - fprintf(stderr, "\n## Numerical Precision Report\n\n"); - fprintf(stderr, "| Component | Operation | max_diff | mean_diff | cosine | Tolerance | Status | Notes |\n"); - fprintf(stderr, "|-----------|-----------|----------|-----------|--------|-----------|--------|-------|\n"); - fprintf(stderr, - "| Tokenizer | Token IDs | %.3e | %.3e | %.8f | exact | %s | %s |\n", - token_row.r.max_diff, - token_row.r.mean_diff, - token_row.r.cosine_sim, - token_row.ok ? "PASS" : "FAIL", - token_row.note.c_str()); - - for (size_t i = 0; i < cases.size(); ++i) { - const auto & tc = cases[i]; - const auto & row = rows[i]; - const char * status = !row.measured ? "UNTESTED" : (row.ok ? "PASS" : "FAIL"); - fprintf(stderr, - "| Phase 5 | %s | %.3e | %.3e | %.8f | %.1e | %s | %s |\n", - tc.label.c_str(), - row.r.max_diff, - row.r.mean_diff, - row.r.cosine_sim, - tolerance_for(tc.name), - status, - row.note.empty() ? "measured" : row.note.c_str()); - } - - return overall_ok ? 0 : 1; -} diff --git a/tests/test_phase6.cpp b/tests/test_phase6.cpp deleted file mode 100644 index 59b1b2a..0000000 --- a/tests/test_phase6.cpp +++ /dev/null @@ -1,287 +0,0 @@ -#include "sam3.h" - -#include "test_utils.h" - -#include -#include -#include -#include - -struct phase6_case { - std::string id; - sam3_pvs_params params; -}; - -struct tensor_case { - std::string name; - std::string label; -}; - -struct report_row { - bool measured = false; - bool ok = false; - compare_result r; - std::string note; -}; - -static std::vector parse_points(const std::string & field) { - std::vector pts; - if (field.empty()) { - return pts; - } - - size_t start = 0; - while (start < field.size()) { - size_t end = field.find('|', start); - if (end == std::string::npos) { - end = field.size(); - } - - const std::string part = field.substr(start, end - start); - size_t colon = part.find(':'); - if (colon != std::string::npos) { - sam3_point pt; - pt.x = std::stof(part.substr(0, colon)); - pt.y = std::stof(part.substr(colon + 1)); - pts.push_back(pt); - } - - start = end + 1; - } - - return pts; -} - -static bool parse_box(const std::string & field, sam3_box & box) { - if (field.empty()) { - return false; - } - - float vals[4]; - size_t start = 0; - for (int i = 0; i < 4; ++i) { - size_t end = field.find(':', start); - if (end == std::string::npos) { - end = field.size(); - } - vals[i] = std::stof(field.substr(start, end - start)); - start = end + 1; - } - - box = {vals[0], vals[1], vals[2], vals[3]}; - return true; -} - -static std::vector load_cases(const std::string & path) { - std::vector cases; - std::ifstream f(path); - std::string line; - while (std::getline(f, line)) { - if (line.empty()) { - continue; - } - - std::vector fields; - size_t start = 0; - while (start <= line.size()) { - size_t end = line.find('\t', start); - if (end == std::string::npos) { - end = line.size(); - } - fields.push_back(line.substr(start, end - start)); - start = end + 1; - if (end == line.size()) { - break; - } - } - while (fields.size() < 5) { - fields.emplace_back(); - } - - phase6_case tc; - tc.id = fields[0]; - tc.params.multimask = fields[1] == "1"; - tc.params.pos_points = parse_points(fields[2]); - tc.params.neg_points = parse_points(fields[3]); - tc.params.use_box = parse_box(fields[4], tc.params.box); - cases.push_back(std::move(tc)); - } - return cases; -} - -static std::vector tensor_cases() { - return { - {"sam_pe_sparse", "Prompt sparse embeddings"}, - {"sam_pe_dense", "Prompt dense embeddings"}, - {"sam_pe_image_pe", "Prompt image PE"}, - {"sam_dec_image_feats", "Decoder image features"}, - {"sam_dec_tokens_initial", "Decoder initial tokens"}, - {"sam_dec_block0_queries", "Decoder block 0 queries"}, - {"sam_dec_block0_keys", "Decoder block 0 keys"}, - {"sam_dec_block1_queries", "Decoder block 1 queries"}, - {"sam_dec_block1_keys", "Decoder block 1 keys"}, - {"sam_dec_final_queries", "Decoder final queries"}, - {"sam_dec_feat_s1_proj", "Decoder s1 projection"}, - {"sam_dec_feat_s0_proj", "Decoder s0 projection"}, - {"sam_dec_upscaled", "Decoder upscaled embedding"}, - {"sam_dec_mask_tokens", "Decoder mask tokens"}, - {"sam_dec_masks", "Decoder mask logits"}, - {"sam_dec_iou", "Decoder IoU predictions"}, - {"sam_dec_obj_score", "Decoder object score logit"}, - {"sam_dec_sam_token", "Decoder SAM token"}, - }; -} - -static float tolerance_for(const std::string & name) { - if (name == "sam_dec_block1_queries" || - name == "sam_dec_final_queries") { - return 2e-4f; - } - if (name == "sam_dec_upscaled") { - return 9e-4f; - } - if (name == "sam_dec_masks") { - return 8e-3f; - } - return 1e-4f; -} - -int main(int argc, char ** argv) { - if (argc < 5) { - fprintf(stderr, "Usage: %s \n", argv[0]); - return 1; - } - - const std::string ref_dir = argv[1]; - const std::string prephase_ref_dir = argv[2]; - const std::string model_path = argv[3]; - const std::string cases_path = argv[4]; - const std::string cpp_root = ref_dir + "/cpp_out_phase6"; - - auto cases = load_cases(cases_path); - if (cases.empty()) { - fprintf(stderr, "No phase 6 cases loaded from %s\n", cases_path.c_str()); - return 1; - } - - if (!ensure_dir(cpp_root)) { - fprintf(stderr, "Failed to create %s\n", cpp_root.c_str()); - return 1; - } - - sam3_params params; - params.model_path = model_path; - params.n_threads = 1; - params.use_gpu = false; - - auto model = sam3_load_model(params); - if (!model) { - fprintf(stderr, "Failed to load model from %s\n", model_path.c_str()); - return 1; - } - - const auto tensors = tensor_cases(); - bool overall_ok = true; - - struct row_print { - std::string case_id; - std::string label; - float max_diff = 0.0f; - float mean_diff = 0.0f; - float cosine = 0.0f; - float tol = 0.0f; - std::string status; - std::string note; - }; - std::vector report; - - for (const auto & tc : cases) { - fprintf(stderr, "\n=== Phase 6 Case %s ===\n", tc.id.c_str()); - - const std::string case_ref_dir = ref_dir + "/" + tc.id; - const std::string case_cpp_dir = cpp_root + "/" + tc.id; - if (!ensure_dir(case_cpp_dir)) { - fprintf(stderr, "Failed to create %s\n", case_cpp_dir.c_str()); - overall_ok = false; - continue; - } - - if (!sam3_test_dump_phase6_from_ref_inputs(*model, prephase_ref_dir, tc.params, - case_cpp_dir, params.n_threads)) { - fprintf(stderr, " [FAIL] could not dump C++ phase 6 tensors for %s\n", tc.id.c_str()); - overall_ok = false; - continue; - } - - for (const auto & tensor : tensors) { - auto ref = load_ref_f32(case_ref_dir + "/" + tensor.name); - auto got = load_ref_f32(case_cpp_dir + "/" + tensor.name); - - row_print row; - row.case_id = tc.id; - row.label = tensor.label; - row.tol = tolerance_for(tensor.name); - - if (ref.data.empty() || got.data.empty()) { - row.status = "FAIL"; - row.note = "missing reference or C++ output"; - overall_ok = false; - report.push_back(std::move(row)); - fprintf(stderr, " [FAIL] missing tensor %s for %s\n", - tensor.name.c_str(), tc.id.c_str()); - continue; - } - - if (ref.numel() != got.numel()) { - row.status = "FAIL"; - row.note = "numel mismatch"; - overall_ok = false; - report.push_back(std::move(row)); - fprintf(stderr, " [FAIL] numel mismatch for %s in %s\n", - tensor.name.c_str(), tc.id.c_str()); - continue; - } - - auto r = compare_tensors(got.data.data(), ref.data.data(), ref.numel(), row.tol); - const bool ok = r.max_diff <= row.tol; - row.max_diff = r.max_diff; - row.mean_diff = r.mean_diff; - row.cosine = r.cosine_sim; - row.status = ok ? "PASS" : "FAIL"; - if (ok) { - row.note = "measured"; - } else { - char note[160]; - snprintf(note, sizeof(note), - "worst_idx=%d got=%.6g ref=%.6g bad=%d/%d", - r.worst_index, r.worst_a, r.worst_b, r.n_bad, r.n_total); - row.note = note; - overall_ok = false; - fprintf(stderr, - " [FAIL] %-24s max=%.3e mean=%.3e cosine=%.8f idx=%d got=%.6g ref=%.6g\n", - tensor.name.c_str(), r.max_diff, r.mean_diff, r.cosine_sim, - r.worst_index, r.worst_a, r.worst_b); - } - - report.push_back(std::move(row)); - } - } - - fprintf(stderr, "\n## Numerical Precision Report\n\n"); - fprintf(stderr, "| Case | Operation | max_diff | mean_diff | cosine | Tolerance | Status | Notes |\n"); - fprintf(stderr, "|------|-----------|----------|-----------|--------|-----------|--------|-------|\n"); - for (const auto & row : report) { - fprintf(stderr, - "| %s | %s | %.3e | %.3e | %.8f | %.1e | %s | %s |\n", - row.case_id.c_str(), - row.label.c_str(), - row.max_diff, - row.mean_diff, - row.cosine, - row.tol, - row.status.c_str(), - row.note.c_str()); - } - - return overall_ok ? 0 : 1; -} diff --git a/tests/test_phase7.cpp b/tests/test_phase7.cpp deleted file mode 100644 index e577495..0000000 --- a/tests/test_phase7.cpp +++ /dev/null @@ -1,244 +0,0 @@ -#include "sam3.h" - -#include "test_utils.h" - -#include -#include -#include -#include -#include - -struct phase7_case { - std::string id; - std::string label; -}; - -struct tensor_case { - std::string name; - std::string label; - float tol; -}; - -struct report_row { - std::string case_id; - std::string label; - float max_diff = 0.0f; - float mean_diff = 0.0f; - float cosine = 0.0f; - float tol = 1e-4f; - std::string status; - std::string note; -}; - -static std::vector load_cases(const std::string & path) { - std::vector cases; - std::ifstream f(path); - std::string line; - while (std::getline(f, line)) { - if (line.empty()) { - continue; - } - std::vector fields; - size_t start = 0; - while (start <= line.size()) { - size_t end = line.find('\t', start); - if (end == std::string::npos) { - end = line.size(); - } - fields.push_back(line.substr(start, end - start)); - start = end + 1; - if (end == line.size()) { - break; - } - } - while (fields.size() < 2) { - fields.emplace_back(); - } - cases.push_back({fields[0], fields[1]}); - } - return cases; -} - -static std::map load_meta(const std::string & path) { - std::map kv; - std::ifstream f(path); - std::string line; - while (std::getline(f, line)) { - if (line.empty()) { - continue; - } - size_t eq = line.find('='); - if (eq == std::string::npos) { - continue; - } - kv[line.substr(0, eq)] = line.substr(eq + 1); - } - return kv; -} - -static std::vector tensor_cases_for(const std::string & case_dir) { - std::vector cases; - const auto meta = load_meta(case_dir + "/meta.txt"); - const int num_slots = meta.count("num_slots") ? std::stoi(meta.at("num_slots")) : 0; - - // Tolerances: pix_proj/obj_ptr/attn_input/sa_l0 are tight (1e-4). - // Memory encoder stages accumulate error through bilinear+conv chain (1e-2). - // Memory attention inherits encoder error non-hermetically (5e-3). - // Propagation accumulates through full pipeline (1e-2). - for (int i = 0; i < num_slots; ++i) { - cases.push_back({"phase7_mem" + std::to_string(i) + "_pix_proj", - "Memory slot " + std::to_string(i) + " pixel projection", 1e-4f}); - cases.push_back({"phase7_mem" + std::to_string(i) + "_fused_input", - "Memory slot " + std::to_string(i) + " fused input", 2e-3f}); - cases.push_back({"phase7_mem" + std::to_string(i) + "_fuser0", - "Memory slot " + std::to_string(i) + " fuser block 0", 1e-2f}); - cases.push_back({"phase7_mem" + std::to_string(i) + "_fuser1", - "Memory slot " + std::to_string(i) + " fuser block 1", 1e-2f}); - cases.push_back({"phase7_mem" + std::to_string(i) + "_output", - "Memory slot " + std::to_string(i) + " encoder output", 1e-2f}); - cases.push_back({"phase7_obj_ptr" + std::to_string(i), - "Memory slot " + std::to_string(i) + " object pointer", 1e-4f}); - } - - cases.push_back({"phase7_mem_attn_input", "Memory attention input", 1e-4f}); - for (int i = 0; i < 4; ++i) { - cases.push_back({"phase7_mem_attn_layer" + std::to_string(i) + "_after_sa", - "Memory attention layer " + std::to_string(i) + " after self-attn", 5e-3f}); - cases.push_back({"phase7_mem_attn_layer" + std::to_string(i) + "_after_ca", - "Memory attention layer " + std::to_string(i) + " after cross-attn", 5e-3f}); - cases.push_back({"phase7_mem_attn_layer" + std::to_string(i) + "_after_ffn", - "Memory attention layer " + std::to_string(i) + " after FFN", 5e-3f}); - } - cases.push_back({"phase7_mem_attn_output", "Memory attention final output", 5e-3f}); - cases.push_back({"phase7_prop_masks", "Propagation mask logits", 1e-2f}); - cases.push_back({"phase7_prop_iou", "Propagation IoU", 1e-3f}); - cases.push_back({"phase7_prop_obj_score", "Propagation object score", 1e-3f}); - cases.push_back({"phase7_prop_sam_token", "Propagation SAM token", 1e-2f}); - return cases; -} - -int main(int argc, char ** argv) { - if (argc < 4) { - fprintf(stderr, "Usage: %s \n", argv[0]); - return 1; - } - - const std::string ref_dir = argv[1]; - const std::string model_path = argv[2]; - const std::string cases_path = argv[3]; - const std::string cpp_root = ref_dir + "/cpp_out_phase7"; - - if (!ensure_dir(cpp_root)) { - fprintf(stderr, "Failed to create %s\n", cpp_root.c_str()); - return 1; - } - - auto cases = load_cases(cases_path); - if (cases.empty()) { - fprintf(stderr, "No phase 7 cases loaded from %s\n", cases_path.c_str()); - return 1; - } - - sam3_params params; - params.model_path = model_path; - params.n_threads = 1; - params.use_gpu = false; - - auto model = sam3_load_model(params); - if (!model) { - fprintf(stderr, "Failed to load model from %s\n", model_path.c_str()); - return 1; - } - - bool overall_ok = true; - std::vector report; - - for (const auto & tc : cases) { - fprintf(stderr, "\n=== Phase 7 Case %s ===\n", tc.id.c_str()); - - const std::string case_ref_dir = ref_dir + "/" + tc.id; - const std::string case_cpp_dir = cpp_root + "/" + tc.id; - if (!ensure_dir(case_cpp_dir)) { - fprintf(stderr, "Failed to create %s\n", case_cpp_dir.c_str()); - overall_ok = false; - continue; - } - - if (!sam3_test_dump_phase7_from_ref_inputs(*model, case_ref_dir, case_cpp_dir, params.n_threads)) { - fprintf(stderr, " [FAIL] could not dump C++ phase 7 tensors for %s\n", tc.id.c_str()); - overall_ok = false; - continue; - } - - for (const auto & tensor : tensor_cases_for(case_ref_dir)) { - auto ref = load_ref_f32(case_ref_dir + "/" + tensor.name); - auto got = load_ref_f32(case_cpp_dir + "/" + tensor.name); - - report_row row; - row.case_id = tc.id; - row.label = tensor.label; - row.tol = tensor.tol; - - if (ref.data.empty() || got.data.empty()) { - row.status = "FAIL"; - row.note = "missing reference or C++ output"; - overall_ok = false; - report.push_back(std::move(row)); - fprintf(stderr, " [FAIL] missing tensor %s for %s\n", - tensor.name.c_str(), tc.id.c_str()); - continue; - } - - if (ref.numel() != got.numel()) { - row.status = "FAIL"; - row.note = "numel mismatch"; - overall_ok = false; - report.push_back(std::move(row)); - fprintf(stderr, " [FAIL] numel mismatch for %s in %s\n", - tensor.name.c_str(), tc.id.c_str()); - continue; - } - - auto r = compare_tensors(got.data.data(), ref.data.data(), ref.numel(), row.tol); - const bool ok = r.max_diff <= row.tol; - row.max_diff = r.max_diff; - row.mean_diff = r.mean_diff; - row.cosine = r.cosine_sim; - row.status = ok ? "PASS" : "FAIL"; - if (ok) { - row.note = "measured"; - } else { - char note[160]; - snprintf(note, sizeof(note), - "worst_idx=%d got=%.6g ref=%.6g bad=%d/%d", - r.worst_index, r.worst_a, r.worst_b, r.n_bad, r.n_total); - row.note = note; - overall_ok = false; - fprintf(stderr, - " [FAIL] %-28s max=%.3e mean=%.3e cosine=%.8f idx=%d got=%.6g ref=%.6g\n", - tensor.name.c_str(), r.max_diff, r.mean_diff, r.cosine_sim, - r.worst_index, r.worst_a, r.worst_b); - } - - report.push_back(std::move(row)); - } - } - - fprintf(stderr, "\n## Numerical Precision Report\n\n"); - fprintf(stderr, "| Case | Operation | max_diff | mean_diff | cosine | Tolerance | Status | Notes |\n"); - fprintf(stderr, "|------|-----------|----------|-----------|--------|-----------|--------|-------|\n"); - for (const auto & row : report) { - fprintf(stderr, - "| %s | %s | %.3e | %.3e | %.8f | %.1e | %s | %s |\n", - row.case_id.c_str(), - row.label.c_str(), - row.max_diff, - row.mean_diff, - row.cosine, - row.tol, - row.status.c_str(), - row.note.c_str()); - } - - return overall_ok ? 0 : 1; -} diff --git a/tests/test_preprocess_compare.py b/tests/test_preprocess_compare.py deleted file mode 100644 index 8746ad6..0000000 --- a/tests/test_preprocess_compare.py +++ /dev/null @@ -1,70 +0,0 @@ -#!/usr/bin/env python3 -"""Quick comparison: does PIL resize match C++ bilinear resize? - -Instead of replicating C++ in Python (slow), we compare: -1. Python torchvision preprocessing output (already saved as ref_phase3/preprocessed) -2. PIL BILINEAR resize → same normalization (what C++ tries to match) - -If PIL BILINEAR matches torchvision, then C++ should too (if C++ matches PIL). -""" -import numpy as np -from PIL import Image -import torch -from torchvision.transforms import v2 - -image_path = "tests/test_random.jpg" -img_size = 1008 - -# Method 1: torchvision pipeline (what dump_phase3_reference.py uses) -img = Image.open(image_path).convert("RGB") -transform = v2.Compose([ - v2.ToDtype(torch.uint8, scale=True), - v2.Resize(size=(img_size, img_size)), - v2.ToDtype(torch.float32, scale=True), - v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), -]) -img_tensor = v2.functional.to_image(img) -py_tv = transform(img_tensor).numpy() # [3, 1008, 1008] - -# Method 2: PIL resize + same normalization -img_pil = img.resize((img_size, img_size), Image.BILINEAR) -pil_arr = np.array(img_pil, dtype=np.float32) # [H, W, 3] -# Normalize same way: (pixel/255.0 - 0.5) / 0.5 -pil_normalized = (pil_arr / 255.0 - 0.5) / 0.5 # [H, W, 3] -py_pil = pil_normalized.transpose(2, 0, 1) # → [3, H, W] - -# Compare -diff = np.abs(py_tv - py_pil) -print(f"torchvision vs PIL BILINEAR:") -print(f" max_diff: {diff.max():.6e}") -print(f" mean_diff: {diff.mean():.6e}") -cos = np.dot(py_tv.flatten(), py_pil.flatten()) / ( - np.linalg.norm(py_tv.flatten()) * np.linalg.norm(py_pil.flatten())) -print(f" cosine: {cos:.10f}") -print(f" pixels differing (>1e-6): {(diff > 1e-6).sum()} / {diff.size}") - -# Method 3: What the actual SAM3 Python code does -# Check what v2.Resize actually uses under the hood -from torchvision.transforms import InterpolationMode -print(f"\n v2.Resize default interpolation: {v2.Resize(size=(100,100)).interpolation}") -print(f" (BILINEAR={InterpolationMode.BILINEAR})") - -# The key question: does v2.Resize use PIL.BILINEAR or torch.nn.functional.interpolate? -# v2.Resize with antialias=True (default since torchvision 0.17) uses different algorithm -# Let's test with antialias=False to match pure bilinear -transform_no_aa = v2.Compose([ - v2.ToDtype(torch.uint8, scale=True), - v2.Resize(size=(img_size, img_size), antialias=False), - v2.ToDtype(torch.float32, scale=True), - v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), -]) -py_no_aa = transform_no_aa(v2.functional.to_image(img)).numpy() - -diff2 = np.abs(py_tv - py_no_aa) -print(f"\ntorchvision (antialias=True) vs (antialias=False):") -print(f" max_diff: {diff2.max():.6e}") -print(f" mean_diff: {diff2.mean():.6e}") - -# If the default has antialias, the C++ simple bilinear will differ -if diff2.max() > 1e-6: - print(" → antialias IS enabled by default, C++ bilinear will differ from Python") diff --git a/tests/test_preprocess_match.cpp b/tests/test_preprocess_match.cpp deleted file mode 100644 index 58cf631..0000000 --- a/tests/test_preprocess_match.cpp +++ /dev/null @@ -1,57 +0,0 @@ -/** - * Test preprocessing match: run sam3_encode_image with C++ preprocessing - * on the same image, then dump outputs for comparison against Python reference. - * - * Usage: - * ./test_preprocess_match - */ -#include "sam3.h" -#include "test_utils.h" - -#include -#include - -int main(int argc, char** argv) { - if (argc < 4) { - fprintf(stderr, "Usage: %s \n", argv[0]); - return 1; - } - - const std::string model_path = argv[1]; - const std::string image_path = argv[2]; - const std::string out_dir = argv[3]; - ensure_dir(out_dir); - - sam3_params params; - params.model_path = model_path; - params.use_gpu = false; - params.n_threads = 4; - - auto model = sam3_load_model(params); - if (!model) { fprintf(stderr, "Failed to load model\n"); return 1; } - - auto state = sam3_create_state(*model, params); - if (!state) { fprintf(stderr, "Failed to create state\n"); return 1; } - - auto img = sam3_load_image(image_path); - if (img.data.empty()) { fprintf(stderr, "Failed to load image\n"); return 1; } - - fprintf(stderr, "Encoding with C++ preprocessing...\n"); - if (!sam3_encode_image(*state, *model, img)) { - fprintf(stderr, "Encoding failed\n"); - return 1; - } - - // Dump outputs - sam3_dump_state_tensor(*state, "vit_output", out_dir + "/vit_output"); - for (int i = 0; i < 4; ++i) { - char name[32]; - snprintf(name, sizeof(name), "neck_det_%d", i); - sam3_dump_state_tensor(*state, name, out_dir + "/" + name); - } - - fprintf(stderr, "Done. Outputs in %s\n", out_dir.c_str()); - state.reset(); - sam3_free_model(*model); - return 0; -} diff --git a/tests/test_pvs_point.cpp b/tests/test_pvs_point.cpp deleted file mode 100644 index c37bdc1..0000000 --- a/tests/test_pvs_point.cpp +++ /dev/null @@ -1,72 +0,0 @@ -/** - * Quick PVS test: single positive point on the llama image. - * Dumps mask to file for visual inspection. - */ -#include "sam3.h" -#include -#include - -int main(int argc, char ** argv) { - if (argc < 3) { - fprintf(stderr, "Usage: %s [x y]\n", argv[0]); - return 1; - } - - const std::string model_path = argv[1]; - const std::string image_path = argv[2]; - float px = argc > 4 ? atof(argv[3]) : 315.0f; - float py = argc > 4 ? atof(argv[4]) : 250.0f; - - sam3_params params; - params.model_path = model_path; - params.use_gpu = false; - params.n_threads = 4; - - fprintf(stderr, "Loading model...\n"); - auto model = sam3_load_model(params); - if (!model) { fprintf(stderr, "Failed to load model\n"); return 1; } - - auto state = sam3_create_state(*model, params); - if (!state) { fprintf(stderr, "Failed to create state\n"); return 1; } - - fprintf(stderr, "Loading image: %s\n", image_path.c_str()); - auto image = sam3_load_image(image_path); - if (image.data.empty()) { fprintf(stderr, "Failed to load image\n"); return 1; } - fprintf(stderr, "Image: %dx%d\n", image.width, image.height); - - fprintf(stderr, "Encoding image...\n"); - if (!sam3_encode_image(*state, *model, image)) { - fprintf(stderr, "Failed to encode image\n"); return 1; - } - - fprintf(stderr, "\n═══ PVS: point at (%.1f, %.1f) ═══\n", px, py); - - sam3_pvs_params pvs; - pvs.pos_points.push_back({px, py}); - pvs.multimask = false; - - auto result = sam3_segment_pvs(*state, *model, pvs); - fprintf(stderr, "Result: %zu detections\n", result.detections.size()); - - for (size_t i = 0; i < result.detections.size(); ++i) { - const auto& d = result.detections[i]; - fprintf(stderr, " det %zu: score=%.4f iou=%.4f obj=%.4f box=[%.1f,%.1f,%.1f,%.1f] mask=%dx%d\n", - i, d.score, d.iou_score, d.mask.obj_score, - d.box.x0, d.box.y0, d.box.x1, d.box.y1, - d.mask.width, d.mask.height); - - std::string mask_path = "/tmp/pvs_mask_" + std::to_string(i) + ".png"; - sam3_save_mask(d.mask, mask_path); - fprintf(stderr, " Saved: %s\n", mask_path.c_str()); - - // Count mask pixels - int n_on = 0; - for (size_t j = 0; j < d.mask.data.size(); ++j) - if (d.mask.data[j] > 0) n_on++; - fprintf(stderr, " Mask: %d/%zu pixels on (%.1f%%)\n", - n_on, d.mask.data.size(), 100.0f * n_on / d.mask.data.size()); - } - - sam3_free_model(*model); - return 0; -} diff --git a/tests/test_random.jpg b/tests/test_random.jpg deleted file mode 100644 index 76884bb..0000000 Binary files a/tests/test_random.jpg and /dev/null differ diff --git a/tests/test_resize_from_pil.py b/tests/test_resize_from_pil.py deleted file mode 100644 index f1d14aa..0000000 --- a/tests/test_resize_from_pil.py +++ /dev/null @@ -1,70 +0,0 @@ -#!/usr/bin/env python3 -"""Feed PIL-decoded pixels through C++'s bilinear resize (replicated in Python -with double precision) and compare against torchvision resize. - -This isolates the resize from the JPEG decoder. -""" -import numpy as np -import torch -from PIL import Image -from torchvision.transforms import v2 - -image_path = "tests/test_random.jpg" -img_size = 1008 - -# Load with PIL (same decoder as torchvision) -img = Image.open(image_path).convert("RGB") -src = np.array(img, dtype=np.uint8) # [480, 640, 3] HWC -src_h, src_w = src.shape[:2] - -# ═══ Torch resize (ground truth) ═══ -img_tensor = v2.functional.to_image(img) -tv = v2.Resize(size=(img_size, img_size))(img_tensor) -tv_hwc = tv.permute(1, 2, 0).numpy() - -# ═══ C++ resize replicated (double precision, same source pixels) ═══ -sx = float(src_w) / img_size -sy = float(src_h) / img_size - -cpp_hwc = np.zeros((img_size, img_size, 3), dtype=np.uint8) -for y in range(img_size): - fy = (y + 0.5) * sy - 0.5 - if fy < 0: fy = 0.0 - y0 = int(fy) - y1 = y0 + 1 if y0 < src_h - 1 else y0 - wy = fy - y0 - wy0 = 1.0 - wy - for x in range(img_size): - fx = (x + 0.5) * sx - 0.5 - if fx < 0: fx = 0.0 - x0 = int(fx) - x1 = x0 + 1 if x0 < src_w - 1 else x0 - wx = fx - x0 - wx0 = 1.0 - wx - for c in range(3): - p00 = float(src[y0, x0, c]) - p01 = float(src[y0, x1, c]) - p10 = float(src[y1, x0, c]) - p11 = float(src[y1, x1, c]) - v = wy0 * (wx0 * p00 + wx * p01) + wy * (wx0 * p10 + wx * p11) - iv = int(v + 0.5) - if iv < 0: iv = 0 - if iv > 255: iv = 255 - cpp_hwc[y, x, c] = iv - -# Compare -diff = np.abs(cpp_hwc.astype(np.int16) - tv_hwc.astype(np.int16)) -print("═══ Same source pixels: C++ resize (double) vs torchvision ═══") -print(f" max pixel diff: {diff.max()}") -print(f" mean pixel diff: {diff.mean():.6f}") -print(f" pixels differing: {(diff > 0).sum()} / {diff.size} ({100*(diff>0).sum()/diff.size:.2f}%)") -for t in [0, 1, 2]: - n = (diff > t).sum() - if n > 0: - print(f" pixels with diff > {t}: {n}") - -if diff.max() == 0: - print("\n PERFECT MATCH! The resize algorithms are identical.") - print(" The 8% difference was entirely from the JPEG decoder (stb_image vs libjpeg).") -else: - print(f"\n {(diff > 0).sum()} pixels still differ — residual rounding difference.") diff --git a/tests/test_sam2_backbone_compare.cpp b/tests/test_sam2_backbone_compare.cpp deleted file mode 100644 index 153f3ae..0000000 --- a/tests/test_sam2_backbone_compare.cpp +++ /dev/null @@ -1,47 +0,0 @@ -// Test: Compare SAM2 backbone + FPN outputs against Python. -// Set SAM2_DUMP_DIR env var to dump debug intermediates. -// Usage: SAM2_DUMP_DIR=/tmp/debug_sam2_cpp test_sam2_backbone_compare - -#include "sam3.h" -#include -#include -#include -#include - -int main(int argc, char** argv) { - if (argc < 4) { - fprintf(stderr, "Usage: %s \n", argv[0]); - return 1; - } - - mkdir(argv[3], 0755); - - std::ifstream fin(argv[2], std::ios::binary); - fin.seekg(0, std::ios::end); - size_t sz = fin.tellg(); - fin.seekg(0); - std::vector img(sz / 4); - fin.read(reinterpret_cast(img.data()), sz); - - sam3_params p; - p.model_path = argv[1]; - p.n_threads = 8; - p.use_gpu = false; - auto model = sam3_load_model(p); - if (!model) return 1; - auto state = sam3_create_state(*model, p); - if (!state) return 1; - - if (!sam3_encode_image_from_preprocessed(*state, *model, img.data(), 1024)) return 1; - - // Dump FPN outputs via state tensor API - for (int i = 0; i < 3; ++i) { - char name[32]; - snprintf(name, sizeof(name), "neck_trk_%d", i); - std::string path = std::string(argv[3]) + "/" + name; - sam3_dump_state_tensor(*state, name, path); - } - - fprintf(stderr, "Done.\n"); - return 0; -} diff --git a/tests/test_sam2_e2e_all_variants.py b/tests/test_sam2_e2e_all_variants.py deleted file mode 100644 index bdff6fc..0000000 --- a/tests/test_sam2_e2e_all_variants.py +++ /dev/null @@ -1,345 +0,0 @@ -#!/usr/bin/env python3 -"""End-to-end comparison of all SAM2.1 variants: Python vs C++ ggml. - -This script: -1. Runs Python SAM2 inference for all 4 variants (Tiny, Small, Base+, Large) -2. Runs C++ ggml inference for the same variants -3. Compares binary masks and IoU scores - -Usage: - cd ~/Documents/sam2 - python /path/to/test_sam2_e2e_all_variants.py \ - --image /path/to/test_image.jpg \ - --cpp-dir /path/to/sam3.cpp-sam2 \ - --point-x 600 --point-y 599 - -Requires: SAM2 Python package installed, C++ test binary built. -""" -import argparse, os, sys, subprocess, json, time -import numpy as np - -sys.path.insert(0, os.path.expanduser("~/Documents/sam2")) -import torch -torch.set_default_dtype(torch.float32) -from PIL import Image - -VARIANTS_21 = [ - {"name": "tiny", "config": "configs/sam2.1/sam2.1_hiera_t.yaml", "prefix": "sam2.1"}, - {"name": "small", "config": "configs/sam2.1/sam2.1_hiera_s.yaml", "prefix": "sam2.1"}, - {"name": "base_plus", "config": "configs/sam2.1/sam2.1_hiera_b+.yaml", "prefix": "sam2.1"}, - {"name": "large", "config": "configs/sam2.1/sam2.1_hiera_l.yaml", "prefix": "sam2.1"}, -] - -VARIANTS_20 = [ - {"name": "tiny", "config": "configs/sam2/sam2_hiera_t.yaml", "prefix": "sam2"}, - {"name": "small", "config": "configs/sam2/sam2_hiera_s.yaml", "prefix": "sam2"}, - {"name": "base_plus", "config": "configs/sam2/sam2_hiera_b+.yaml", "prefix": "sam2"}, - {"name": "large", "config": "configs/sam2/sam2_hiera_l.yaml", "prefix": "sam2"}, -] - - -def run_python_variant(variant_name, config, checkpoint, image_path, point_x, point_y, dump_dir): - """Run Python SAM2 inference for one variant.""" - from sam2.build_sam import build_sam2 - from sam2.sam2_image_predictor import SAM2ImagePredictor - - os.makedirs(dump_dir, exist_ok=True) - - print(f"\n{'='*60}") - print(f"Python: {variant_name}") - print(f"{'='*60}") - - t0 = time.time() - model = build_sam2(config, checkpoint, device="cpu") - model = model.float().eval() - predictor = SAM2ImagePredictor(model) - t_load = time.time() - t0 - - image = np.array(Image.open(image_path).convert("RGB")) - img_h, img_w = image.shape[:2] - - t0 = time.time() - with torch.no_grad(): - predictor.set_image(image) - t_encode = time.time() - t0 - - point_coords = np.array([[point_x, point_y]], dtype=np.float32) - point_labels = np.array([1], dtype=np.int32) - - t0 = time.time() - with torch.no_grad(): - masks, scores, logits = predictor.predict( - point_coords=point_coords, - point_labels=point_labels, - multimask_output=True, - normalize_coords=True, - ) - t_predict = time.time() - t0 - - # Save preprocessed image for C++ consumption - transforms = predictor._transforms - img_tensor = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0 - from torchvision.transforms.functional import resize, normalize - img_resized = resize(img_tensor, [model.image_size, model.image_size], antialias=True) - mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1) - std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1) - img_normalized = (img_resized - mean) / std - preproc_path = os.path.join(dump_dir, "preprocessed.bin") - img_normalized.numpy().astype(np.float32).tofile(preproc_path) - - # Save masks and scores - logits.astype(np.float32).tofile(os.path.join(dump_dir, "logits.bin")) - scores.astype(np.float32).tofile(os.path.join(dump_dir, "scores.bin")) - masks.astype(np.uint8).tofile(os.path.join(dump_dir, "masks.bin")) - - # Save metadata - meta = { - "variant": variant_name, "img_w": img_w, "img_h": img_h, - "point_x": point_x, "point_y": point_y, - "n_masks": int(masks.shape[0]), - "mask_h": int(masks.shape[1]), "mask_w": int(masks.shape[2]), - "logit_h": int(logits.shape[1]), "logit_w": int(logits.shape[2]), - "scores": scores.tolist(), - "t_load": t_load, "t_encode": t_encode, "t_predict": t_predict, - } - with open(os.path.join(dump_dir, "meta.json"), "w") as f: - json.dump(meta, f, indent=2) - - best = int(np.argmax(scores)) - fg_pct = 100 * masks[best].sum() / masks[best].size - print(f" Load: {t_load:.1f}s Encode: {t_encode:.1f}s Predict: {t_predict:.3f}s") - print(f" Scores: {scores}") - print(f" Best mask {best}: {fg_pct:.1f}% foreground") - - return meta - - -def run_cpp_variant(variant_name, model_path, preproc_path, dump_dir, cpp_binary, point_x, point_y, img_w, img_h): - """Run C++ ggml inference for one variant.""" - print(f"\n C++: {variant_name}") - - os.makedirs(dump_dir, exist_ok=True) - - # Build a minimal C++ test that uses sam3_encode_image_from_preprocessed + sam3_segment_pvs - # and writes raw mask logits + IoU scores - env = os.environ.copy() - env["SAM2_DUMP_DIR"] = dump_dir - build_dir = os.path.dirname(cpp_binary) - env["DYLD_LIBRARY_PATH"] = ":".join([ - os.path.join(build_dir, "ggml/src"), - os.path.join(build_dir, "ggml/src/ggml-cpu"), - os.path.join(build_dir, "ggml/src/ggml-metal"), - os.path.join(build_dir, "ggml/src/ggml-blas"), - ]) - - cmd = [cpp_binary, model_path, preproc_path, dump_dir, - str(img_w), str(img_h), str(point_x), str(point_y)] - t0 = time.time() - result = subprocess.run(cmd, capture_output=True, text=True, env=env, timeout=600) - t_total = time.time() - t0 - - if result.returncode != 0: - print(f" C++ FAILED (exit {result.returncode}):") - print(result.stderr[-500:] if result.stderr else "no stderr") - return None - - # Parse IoU from stderr - ious = [] - for line in result.stderr.split("\n"): - if "iou=" in line: - # Parse: iou=[0.990, 0.994, 0.397, 0.977] - import re - m = re.search(r'iou=\[([\d., ]+)\]', line) - if m: - ious = [float(x.strip()) for x in m.group(1).split(",")] - - print(f" Time: {t_total:.1f}s IoU: {ious}") - return {"ious": ious, "t_total": t_total} - - -def compare_variant(variant_name, py_dir, cpp_dir, py_meta): - """Compare Python and C++ results for one variant.""" - print(f"\n Comparison: {variant_name}") - - # Load Python masks and scores - py_scores = np.fromfile(os.path.join(py_dir, "scores.bin"), dtype=np.float32) - py_logits = np.fromfile(os.path.join(py_dir, "logits.bin"), dtype=np.float32) - n_masks = py_meta["n_masks"] - logit_h, logit_w = py_meta["logit_h"], py_meta["logit_w"] - py_logits = py_logits.reshape(n_masks, logit_h, logit_w) - - # Load C++ mask logits: [H4*H4, 4, 1, 1] in ggml layout - cpp_mask_path = os.path.join(cpp_dir, "cpp_pvs_masks.bin") - if not os.path.exists(cpp_mask_path): - print(f" SKIP — C++ masks not found") - return None - - cpp_masks_raw = np.fromfile(cpp_mask_path, dtype=np.float32) - # C++ has 4 masks (including single-mask at index 0), Python has 3 (multimask only) - # C++ ggml layout: [H4*H4, 4, 1, 1] where flat = pixel + mask*H4*H4 - # Figure out mask resolution from total size - n_pixels = cpp_masks_raw.size // 4 - mask_hw = int(np.sqrt(n_pixels)) - cpp_masks_4 = cpp_masks_raw.reshape(4, n_pixels) - - # C++ IoU - cpp_iou_path = os.path.join(cpp_dir, "cpp_pvs_iou.bin") - cpp_ious = np.fromfile(cpp_iou_path, dtype=np.float32) if os.path.exists(cpp_iou_path) else np.zeros(4) - - # Compare multimask outputs (indices 1-3 in C++, 0-2 in Python) - results = [] - for m in range(n_masks): - py_m = py_logits[m].flatten() - cpp_m = cpp_masks_4[m + 1] # skip index 0 (single-mask) - - # Resize if needed (Python is logit_h × logit_w, C++ is mask_hw × mask_hw) - if py_m.size != cpp_m.size: - # Reshape and compare at the smaller resolution - print(f" Mask {m}: SIZE MISMATCH py={py_m.size} cpp={cpp_m.size}") - results.append({"mask": m, "status": "SIZE_MISMATCH"}) - continue - - # Cosine similarity of logits - cos = np.dot(py_m, cpp_m) / (np.linalg.norm(py_m) * np.linalg.norm(cpp_m) + 1e-8) - - # Binary mask agreement (threshold at 0) - py_bin = (py_m > 0).astype(int) - cpp_bin = (cpp_m > 0).astype(int) - union = max((py_bin | cpp_bin).sum(), 1) - intersect = (py_bin & cpp_bin).sum() - binary_iou = intersect / union - agree_pct = 100 * (py_bin == cpp_bin).sum() / len(py_bin) - - # IoU score comparison - py_iou = py_scores[m] - cpp_iou = cpp_ious[m + 1] if len(cpp_ious) > m + 1 else 0 - - status = "PASS" if binary_iou > 0.90 else "FAIL" - results.append({ - "mask": m, "status": status, - "logit_cos": float(cos), "binary_iou": float(binary_iou), - "agree_pct": float(agree_pct), - "py_iou_score": float(py_iou), "cpp_iou_score": float(cpp_iou), - "py_fg": int(py_bin.sum()), "cpp_fg": int(cpp_bin.sum()), - }) - - flag = "✓" if status == "PASS" else "✗" - print(f" {flag} Mask {m}: logit_cos={cos:.4f} binary_IoU={binary_iou:.4f} " - f"agree={agree_pct:.1f}% py_iou={py_iou:.3f} cpp_iou={cpp_iou:.3f}") - - return results - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--image", required=True) - parser.add_argument("--cpp-dir", required=True, help="Path to sam3.cpp-sam2 root") - parser.add_argument("--point-x", type=float, default=600) - parser.add_argument("--point-y", type=float, default=599) - parser.add_argument("--variants", default="tiny,small,base_plus,large", - help="Comma-separated list of variants to test") - parser.add_argument("--version", default="2.1", choices=["2.0", "2.1", "both"], - help="SAM2 version to test") - args = parser.parse_args() - - cpp_dir = args.cpp_dir - cpp_binary = os.path.join(cpp_dir, "build", "test_sam2_pvs_compare") - if not os.path.exists(cpp_binary): - print(f"ERROR: C++ binary not found at {cpp_binary}") - print(f"Build it first: cd {cpp_dir}/build && make test_sam2_pvs_compare") - sys.exit(1) - - selected = args.variants.split(",") - base_dump = "/tmp/sam2_e2e_test" - os.makedirs(base_dump, exist_ok=True) - - image = np.array(Image.open(args.image).convert("RGB")) - img_h, img_w = image.shape[:2] - - all_results = {} - - # Select variant list based on version - if args.version == "both": - variant_lists = [("SAM2.0", VARIANTS_20), ("SAM2.1", VARIANTS_21)] - elif args.version == "2.0": - variant_lists = [("SAM2.0", VARIANTS_20)] - else: - variant_lists = [("SAM2.1", VARIANTS_21)] - - for ver_label, variants in variant_lists: - print(f"\n{'#'*70}") - print(f"# {ver_label}") - print(f"{'#'*70}") - - for v in variants: - if v["name"] not in selected: - continue - - name = v["name"] - prefix = v["prefix"] - display_name = f"{prefix}_{name}" - ckpt = os.path.join(cpp_dir, f"weights/{prefix}/{prefix}_hiera_{name}.pt") - ggml = os.path.join(cpp_dir, f"weights/ggml/{prefix}_hiera_{name}_f32.ggml") - config = v["config"] - - if not os.path.exists(ckpt): - print(f"\nSKIP {display_name}: checkpoint not found at {ckpt}") - continue - if not os.path.exists(ggml): - print(f"\nSKIP {display_name}: ggml model not found at {ggml}") - continue - - py_dump = os.path.join(base_dump, f"py_{display_name}") - cpp_dump = os.path.join(base_dump, f"cpp_{display_name}") - - # Run Python - py_meta = run_python_variant( - display_name, config, ckpt, args.image, - args.point_x, args.point_y, py_dump - ) - - # Run C++ - cpp_result = run_cpp_variant( - display_name, ggml, - os.path.join(py_dump, "preprocessed.bin"), - cpp_dump, cpp_binary, - args.point_x, args.point_y, img_w, img_h - ) - - # Compare - if cpp_result is not None: - results = compare_variant(display_name, py_dump, cpp_dump, py_meta) - all_results[display_name] = results - - # Summary - print(f"\n{'='*70}") - print(f"SUMMARY: SAM2.1 End-to-End Comparison (all variants)") - print(f"{'='*70}") - print(f"Image: {args.image} ({img_w}x{img_h})") - print(f"Point: ({args.point_x}, {args.point_y})") - print() - print(f"{'Variant':<12} {'Best Mask IoU':>14} {'Logit Cos':>10} {'Binary IoU':>11} {'Agreement':>10} {'Status':>8}") - print("-" * 70) - - all_pass = True - for name, results in all_results.items(): - if not results: - print(f"{name:<12} {'ERROR':>14}") - all_pass = False - continue - best_idx = max(range(len(results)), key=lambda i: results[i].get("py_iou_score", 0)) - r = results[best_idx] - status = r["status"] - if status != "PASS": - all_pass = False - print(f"{name:<12} {r.get('py_iou_score', 0):>14.4f} {r.get('logit_cos', 0):>10.4f} " - f"{r.get('binary_iou', 0):>11.4f} {r.get('agree_pct', 0):>9.1f}% " - f"{' ✓' if status == 'PASS' else ' ✗':>8}") - - print() - print(f"Overall: {'ALL PASS ✓' if all_pass else 'SOME FAILED ✗'}") - return 0 if all_pass else 1 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/tests/test_sam2_pe_cache_size.cpp b/tests/test_sam2_pe_cache_size.cpp deleted file mode 100644 index d01672d..0000000 --- a/tests/test_sam2_pe_cache_size.cpp +++ /dev/null @@ -1,116 +0,0 @@ -// Regression test: PE cache must use feat_size() (64 for SAM2), NOT n_img_embd() (73). -// -// The bug: sam3_populate_pe_cache() used n_img_embd() = img_size/patch_size = 1024/14 = 73 -// for the dense PE grid size. For SAM2 (Hiera backbone), the correct size is feat_size() = 64. -// A 73×73 PE grid fed to a 64×64 decoder produces garbage masks. -// -// This test verifies: -// 1. The PE cache dense grids have size D * feat_size^2 (not D * n_img_embd^2) -// 2. The PVS decoder produces a mask where the clicked point is inside the foreground -// 3. The IoU score is above a reasonable threshold -// -// Usage: test_sam2_pe_cache_size - -#include "sam3.h" -#include -#include -#include - -int main(int argc, char** argv) { - if (argc < 3) { - fprintf(stderr, "Usage: %s \n", argv[0]); - return 1; - } - - sam3_params p; - p.model_path = argv[1]; - p.n_threads = 4; - p.use_gpu = false; - - auto model = sam3_load_model(p); - if (!model) { fprintf(stderr, "FAIL: cannot load model\n"); return 1; } - - // Must be a SAM2 model - auto mt = sam3_get_model_type(*model); - if (mt != SAM3_MODEL_SAM2) { - fprintf(stderr, "SKIP: not a SAM2 model (type=%d)\n", mt); - return 0; - } - - auto state = sam3_create_state(*model, p); - if (!state) { fprintf(stderr, "FAIL: cannot create state\n"); return 1; } - - // Load and encode image - auto image = sam3_load_image(argv[2]); - if (image.data.empty()) { fprintf(stderr, "FAIL: cannot load image\n"); return 1; } - fprintf(stderr, "Image: %dx%d\n", image.width, image.height); - - if (!sam3_encode_image(*state, *model, image)) { - fprintf(stderr, "FAIL: image encoding failed\n"); - return 1; - } - - // Run PVS with a point at the center of the image - float cx = image.width / 2.0f; - float cy = image.height / 2.0f; - fprintf(stderr, "Point: (%.0f, %.0f)\n", cx, cy); - - sam3_pvs_params pvs; - pvs.pos_points.push_back({cx, cy}); - pvs.multimask = false; // single mask — same as GUI default - - auto result = sam3_segment_pvs(*state, *model, pvs); - - // ── Test 1: Must produce exactly 1 detection ── - if (result.detections.empty()) { - fprintf(stderr, "FAIL: no detections returned\n"); - return 1; - } - fprintf(stderr, "Detections: %zu\n", result.detections.size()); - - auto& det = result.detections[0]; - fprintf(stderr, "Mask 0: iou=%.4f, mask=%dx%d\n", - det.iou_score, det.mask.width, det.mask.height); - - // ── Test 2: IoU score must be reasonable (> 0.1) ── - // A garbage PE grid produces IoU near 0 or random values - if (det.iou_score < 0.1f) { - fprintf(stderr, "FAIL: IoU score %.4f is too low (PE cache likely wrong size)\n", - det.iou_score); - return 1; - } - - // ── Test 3: The clicked point must be INSIDE the foreground mask ── - // If the PE grid is wrong size, the mask is often inverted or displaced - int px = (int)(cx * det.mask.width / image.width); - int py = (int)(cy * det.mask.height / image.height); - if (px >= det.mask.width) px = det.mask.width - 1; - if (py >= det.mask.height) py = det.mask.height - 1; - int idx = py * det.mask.width + px; - bool point_in_mask = (det.mask.data[idx] > 0); - - if (!point_in_mask) { - fprintf(stderr, "FAIL: clicked point (%d,%d) is NOT inside the predicted mask\n", px, py); - fprintf(stderr, " This indicates the PE cache spatial size is wrong.\n"); - fprintf(stderr, " Expected feat_size()=%d, check sam3_populate_pe_cache().\n", - 64); // SAM2 feat_size - return 1; - } - - // ── Test 4: Foreground coverage must be reasonable (0.1% - 99%) ── - // A wrong PE grid often produces all-foreground or all-background masks - int fg_count = 0; - for (size_t i = 0; i < det.mask.data.size(); ++i) - if (det.mask.data[i] > 0) fg_count++; - float fg_pct = 100.0f * fg_count / det.mask.data.size(); - - if (fg_pct < 0.05f || fg_pct > 99.5f) { - fprintf(stderr, "FAIL: foreground coverage %.2f%% is degenerate (expected 0.05-99.5%%)\n", - fg_pct); - return 1; - } - - fprintf(stderr, "PASS: IoU=%.4f, point_in_mask=true, fg=%.1f%%\n", - det.iou_score, fg_pct); - return 0; -} diff --git a/tests/test_sam2_pvs_compare.cpp b/tests/test_sam2_pvs_compare.cpp deleted file mode 100644 index 93efbd6..0000000 --- a/tests/test_sam2_pvs_compare.cpp +++ /dev/null @@ -1,74 +0,0 @@ -// Test: Compare SAM2 PVS decoder output against Python reference. -// Set SAM2_DUMP_DIR to dump debug tensors. -// Usage: test_sam2_pvs_compare [orig_w orig_h point_x point_y] -// Default: orig 1200x1198, point at (600, 599) - -#include "sam3.h" -#include -#include -#include -#include -#include - -// Hack: set orig dims on state (state is opaque, but we know its layout) -// This is needed because encode_image_from_preprocessed sets orig to img_size -extern void sam3_state_set_orig_dims(sam3_state& state, int w, int h); - -int main(int argc, char** argv) { - if (argc < 4) { - fprintf(stderr, "Usage: %s [orig_w orig_h point_x point_y]\n", argv[0]); - return 1; - } - - int orig_w = (argc > 4) ? atoi(argv[4]) : 1200; - int orig_h = (argc > 5) ? atoi(argv[5]) : 1198; - float point_x = (argc > 6) ? atof(argv[6]) : 600.0f; - float point_y = (argc > 7) ? atof(argv[7]) : 599.0f; - - mkdir(argv[3], 0755); - - std::ifstream fin(argv[2], std::ios::binary); - fin.seekg(0, std::ios::end); - size_t sz = fin.tellg(); - fin.seekg(0); - std::vector img(sz / 4); - fin.read(reinterpret_cast(img.data()), sz); - - sam3_params p; - p.model_path = argv[1]; - p.n_threads = 8; - p.use_gpu = false; - auto model = sam3_load_model(p); - if (!model) return 1; - auto state = sam3_create_state(*model, p); - if (!state) return 1; - - if (!sam3_encode_image_from_preprocessed(*state, *model, img.data(), 1024)) return 1; - - // Override original image dimensions (encode_from_preprocessed sets them to 1024) - sam3_state_set_orig_dims(*state, orig_w, orig_h); - - sam3_pvs_params pvs; - pvs.pos_points.push_back({point_x, point_y}); - pvs.multimask = true; - - auto result = sam3_segment_pvs(*state, *model, pvs); - - fprintf(stderr, "PVS result: %zu detections\n", result.detections.size()); - for (size_t i = 0; i < result.detections.size(); ++i) { - fprintf(stderr, " det %zu: iou=%.4f score=%.4f mask=%dx%d\n", - i, result.detections[i].iou_score, result.detections[i].score, - result.detections[i].mask.width, result.detections[i].mask.height); - } - - if (result.detections.size() > 0) { - for (size_t i = 0; i < result.detections.size(); ++i) { - char path[256]; - snprintf(path, sizeof(path), "%s/mask_%zu.png", argv[3], i); - sam3_save_mask(result.detections[i].mask, path); - } - } - - fprintf(stderr, "Done.\n"); - return 0; -} diff --git a/tests/test_sam2_video_debug.cpp b/tests/test_sam2_video_debug.cpp deleted file mode 100644 index 13917e7..0000000 --- a/tests/test_sam2_video_debug.cpp +++ /dev/null @@ -1,176 +0,0 @@ -/** - * test_sam2_video_debug — Debug SAM2 video tracking pipeline. - * - * Loads bedroom video frames, runs video tracking for 5 frames, - * and dumps intermediate tensors for comparison with Python reference. - * - * Usage: - * test_sam2_video_debug [point_x point_y] - * - * Example: - * test_sam2_video_debug models/sam2.1_hiera_tiny_f32.ggml \ - * ~/Documents/sam2/notebooks/videos/bedroom 210 350 - */ -#include "sam3.h" -#include -#include -#include -#include -#include -#include -#include - -static void dump_tensor_f32(const char* dir, const char* name, - const float* data, int n, - const char* shape_str) { - char path[512]; - snprintf(path, sizeof(path), "%s/%s.bin", dir, name); - FILE* f = fopen(path, "wb"); - if (!f) { fprintf(stderr, "Failed to write %s\n", path); return; } - fwrite(data, sizeof(float), n, f); - fclose(f); - - snprintf(path, sizeof(path), "%s/%s.shape", dir, name); - f = fopen(path, "w"); - if (!f) return; - fprintf(f, "%s", shape_str); - fclose(f); - - // Print stats - float mn = data[0], mx = data[0]; - double sum = 0; - for (int i = 0; i < n; ++i) { - if (data[i] < mn) mn = data[i]; - if (data[i] > mx) mx = data[i]; - sum += data[i]; - } - fprintf(stderr, " [DUMP] %-50s shape=%-20s range=[%.4f, %.4f] mean=%.6f\n", - name, shape_str, mn, mx, sum / n); -} - -int main(int argc, char** argv) { - if (argc < 3) { - fprintf(stderr, "Usage: %s [point_x point_y]\n", argv[0]); - return 1; - } - - const char* model_path = argv[1]; - const char* frames_dir = argv[2]; - float point_x = (argc > 3) ? atof(argv[3]) : 210.0f; - float point_y = (argc > 4) ? atof(argv[4]) : 350.0f; - - const char* dump_dir = "/tmp/debug_sam2_cpp"; - mkdir(dump_dir, 0755); - - const int N_FRAMES = argc > 5 ? atoi(argv[5]) : 5; - int encode_img_size = (argc > 6) ? atoi(argv[6]) : 0; - - // Load model - fprintf(stderr, "Loading model: %s\n", model_path); - sam3_params params; - params.model_path = model_path; - params.use_gpu = false; // CPU for reproducibility - params.encode_img_size = encode_img_size; - auto model = sam3_load_model(params); - if (!model) { - fprintf(stderr, "Failed to load model\n"); - return 1; - } - - // Create state - auto state = sam3_create_state(*model, params); - if (!state) { - fprintf(stderr, "Failed to create state\n"); - return 1; - } - - // Load frames as JPEG images - std::vector frames; - for (int i = 0; i < N_FRAMES + 1; ++i) { - char path[512]; - snprintf(path, sizeof(path), "%s/%05d.jpg", frames_dir, i); - sam3_image img = sam3_load_image(path); - if (img.data.empty()) { - fprintf(stderr, "Failed to load frame %d: %s\n", i, path); - return 1; - } - fprintf(stderr, "Loaded frame %d: %dx%d\n", i, img.width, img.height); - frames.push_back(std::move(img)); - } - - // Create visual tracker - sam3_visual_track_params vtp; - vtp.max_keep_alive = 100; - vtp.recondition_every = 16; - auto tracker = sam3_create_visual_tracker(*model, vtp); - if (!tracker) { - fprintf(stderr, "Failed to create tracker\n"); - return 1; - } - - // ── Frame 0: Encode + add instance with point prompt ────────────────── - fprintf(stderr, "\n=== Frame 0: Encode + Add Instance ===\n"); - - // Encode frame 0 - if (!sam3_encode_image(*state, *model, frames[0])) { - fprintf(stderr, "Failed to encode frame 0\n"); - return 1; - } - - // Dump neck features for frame 0 - // state.neck_trk[0,1,2] are the FPN features - sam3_dump_state_tensor(*state, "neck_trk_0", std::string(dump_dir) + "/f0_neck_trk_0"); - sam3_dump_state_tensor(*state, "neck_trk_1", std::string(dump_dir) + "/f0_neck_trk_1"); - sam3_dump_state_tensor(*state, "neck_trk_2", std::string(dump_dir) + "/f0_neck_trk_2"); - - // Add instance with point prompt - sam3_pvs_params pvs; - pvs.pos_points.push_back({point_x, point_y}); - pvs.multimask = false; - - int inst_id = sam3_tracker_add_instance(*tracker, *state, *model, pvs); - fprintf(stderr, "Added instance: %d\n", inst_id); - if (inst_id < 0) { - fprintf(stderr, "Failed to add instance\n"); - return 1; - } - - // ── Propagate frames 1-4 ────────────────────────────────────────────── - for (int fi = 1; fi <= N_FRAMES; ++fi) { - fprintf(stderr, "\n=== Frame %d: Propagate ===\n", fi); - - auto result = sam3_propagate_frame(*tracker, *state, *model, frames[fi]); - - fprintf(stderr, " Detections: %zu\n", result.detections.size()); - for (size_t d = 0; d < result.detections.size(); ++d) { - const auto& det = result.detections[d]; - int fg = 0; - for (auto v : det.mask.data) if (v > 127) fg++; - float pct = 100.0f * fg / std::max(1, (int)det.mask.data.size()); - fprintf(stderr, " det[%zu] inst=%d score=%.4f fg=%d (%.1f%%) mask=%dx%d\n", - d, det.instance_id, det.score, fg, pct, - det.mask.width, det.mask.height); - } - - // Dump neck features - char buf[128]; - snprintf(buf, sizeof(buf), "f%d_neck_trk_2", fi); - sam3_dump_state_tensor(*state, "neck_trk_2", - std::string(dump_dir) + "/" + buf); - - // Dump output mask - if (!result.detections.empty()) { - const auto& det = result.detections[0]; - std::vector mask_f(det.mask.data.size()); - for (size_t i = 0; i < det.mask.data.size(); ++i) - mask_f[i] = det.mask.data[i] > 127 ? 1.0f : 0.0f; - snprintf(buf, sizeof(buf), "f%d_output_mask_binary", fi); - char shape[64]; - snprintf(shape, sizeof(shape), "%d,%d", det.mask.height, det.mask.width); - dump_tensor_f32(dump_dir, buf, mask_f.data(), (int)mask_f.size(), shape); - } - } - - fprintf(stderr, "\nDone. Tensors dumped to %s/\n", dump_dir); - return 0; -} diff --git a/tests/test_text_enc_cat.cpp b/tests/test_text_enc_cat.cpp deleted file mode 100644 index 760be33..0000000 --- a/tests/test_text_enc_cat.cpp +++ /dev/null @@ -1,75 +0,0 @@ -/** - * Quick test: run the text encoder in isolation for "cat" token IDs - * and compare against Python reference. - */ -#include "sam3.h" -#include "test_utils.h" -#include -#include -#include - -int main(int argc, char ** argv) { - if (argc < 3) { - fprintf(stderr, "Usage: %s \n", argv[0]); - return 1; - } - - const std::string model_path = argv[1]; - const std::string ref_dir = std::string(argv[2]) + "/pcs_ref"; - const std::string out_dir = std::string(argv[2]) + "/pcs_cpp_text_only"; - ensure_dir(out_dir); - - sam3_params params; - params.model_path = model_path; - params.use_gpu = false; - params.n_threads = 4; - - auto model = sam3_load_model(params); - if (!model) { fprintf(stderr, "Failed to load model\n"); return 1; } - - // Load cat token IDs from Python reference - auto ref_tokens_f = load_ref_f32(ref_dir + "/token_ids"); - if (ref_tokens_f.data.empty()) { - fprintf(stderr, "Failed to load token IDs\n"); - return 1; - } - std::vector token_ids(ref_tokens_f.numel()); - for (int i = 0; i < ref_tokens_f.numel(); ++i) { - token_ids[i] = (int32_t)ref_tokens_f.data[i]; - } - while ((int)token_ids.size() < 32) token_ids.push_back(0); - - fprintf(stderr, "Cat token IDs: [%d, %d, %d, ...]\n", - token_ids[0], token_ids[1], token_ids[2]); - - // Run text encoder in ISOLATION - if (!sam3_test_dump_text_encoder(*model, token_ids, out_dir, params.n_threads)) { - fprintf(stderr, "Failed to dump text encoder\n"); - return 1; - } - - // Compare against Python reference - auto cpp = load_ref_f32(out_dir + "/text_features_2d"); - auto py = load_ref_f32(ref_dir + "/text_features"); - if (cpp.data.empty() || py.data.empty()) { - fprintf(stderr, "Failed to load tensors for comparison\n"); - return 1; - } - - int n = std::min(cpp.numel(), py.numel()); - auto r = compare_tensors(cpp.data.data(), py.data.data(), n, 1e-4f); - - fprintf(stderr, "\n═══ Text Encoder (Isolated): C++ vs Python for 'cat' ═══\n"); - fprintf(stderr, " shape: cpp=[%d,%d] py=[%d,%d]\n", - cpp.shape[0], cpp.shape.size()>1?cpp.shape[1]:0, - py.shape[0], py.shape.size()>1?py.shape[1]:0); - fprintf(stderr, " max_diff: %.6e\n", r.max_diff); - fprintf(stderr, " mean_diff: %.6e\n", r.mean_diff); - fprintf(stderr, " cosine: %.8f\n", r.cosine_sim); - fprintf(stderr, " n_bad (atol=1e-4): %d/%d\n", r.n_bad, r.n_total); - - fprintf(stderr, "\n Result: %s\n", r.cosine_sim > 0.999f ? "PASS" : "FAIL"); - - sam3_free_model(*model); - return r.cosine_sim > 0.999f ? 0 : 1; -} diff --git a/tests/test_text_enc_dump.cpp b/tests/test_text_enc_dump.cpp deleted file mode 100644 index 174d647..0000000 --- a/tests/test_text_enc_dump.cpp +++ /dev/null @@ -1,58 +0,0 @@ -#include "sam3.h" - -#include -#include -#include -#include - -int main(int argc, char ** argv) { - if (argc < 3) { - fprintf(stderr, "Usage: %s \n", argv[0]); - return 1; - } - - const std::string model_path = argv[1]; - const std::string output_dir = argv[2]; - - // Create output directory - mkdir(output_dir.c_str(), 0755); - - // Load tokenizer from embedded data in model file - if (!sam3_test_load_tokenizer(model_path)) { - fprintf(stderr, "Failed to load tokenizer from %s\n", model_path.c_str()); - return 1; - } - - // Tokenize "shoe" - auto token_ids = sam3_test_tokenize("shoe"); - fprintf(stderr, "Token IDs: ["); - for (size_t i = 0; i < token_ids.size(); ++i) { - if (i > 0) fprintf(stderr, ", "); - fprintf(stderr, "%d", token_ids[i]); - if (token_ids[i] == 0 && i > 1) { fprintf(stderr, ", ..."); break; } - } - fprintf(stderr, "]\n"); - - // Load model - sam3_params params; - params.model_path = model_path; - params.n_threads = 1; - params.use_gpu = false; - - auto model = sam3_load_model(params); - if (!model) { - fprintf(stderr, "Failed to load model from %s\n", model_path.c_str()); - return 1; - } - - // Dump text encoder tensors - if (!sam3_test_dump_text_encoder(*model, token_ids, output_dir, params.n_threads)) { - fprintf(stderr, "Failed to dump text encoder tensors\n"); - sam3_free_model(*model); - return 1; - } - - fprintf(stderr, "All tensors dumped to %s/\n", output_dir.c_str()); - sam3_free_model(*model); - return 0; -} diff --git a/tests/test_tokenizer.cpp b/tests/test_tokenizer.cpp deleted file mode 100644 index 223637b..0000000 --- a/tests/test_tokenizer.cpp +++ /dev/null @@ -1,87 +0,0 @@ -#include "sam3.h" - -#include -#include -#include -#include -#include - -// ── Test cases with Python-verified expected token IDs ─────────────────────── -// Generated by: python scripts/verify_tokenizer.py raw_weights/ - -struct test_case { - const char * text; - std::vector expected; // non-padding tokens only (SOT + BPE + EOT) -}; - -static const test_case CASES[] = { - { "a", { 49406, 320, 49407 } }, - { "hello world", { 49406, 3306, 1002, 49407 } }, - { "yellow school bus", { 49406, 4481, 1228, 2840, 49407 } }, - { "a photo of a cat", { 49406, 320, 1125, 539, 320, 2368, 49407 } }, - { "don't stop", { 49406, 847, 713, 1691, 49407 } }, - { "SAM 3", { 49406, 3730, 274, 49407 } }, - { "123", { 49406, 272, 273, 274, 49407 } }, - { "", { 49406, 49407 } }, - { "a person riding a bicycle on a sunny day in the park", - { 49406, 320, 2533, 6765, 320, 11652, 525, 320, 5438, 575, 530, 518, 1452, 49407 } }, -}; -static const int N_CASES = sizeof(CASES) / sizeof(CASES[0]); - -int main(int argc, char ** argv) { - if (argc < 2) { - fprintf(stderr, "Usage: %s \n", argv[0]); - return 1; - } - - if (!sam3_test_load_tokenizer(argv[1])) { - fprintf(stderr, "FAIL: could not load tokenizer from '%s'\n", argv[1]); - return 1; - } - - int n_pass = 0; - for (int i = 0; i < N_CASES; ++i) { - auto tokens = sam3_test_tokenize(CASES[i].text); - const auto & exp = CASES[i].expected; - bool pass = true; - - // Must be exactly 32 tokens (context length) - if ((int)tokens.size() != 32) { - fprintf(stderr, "FAIL [%d]: expected 32 tokens, got %d\n", i, (int)tokens.size()); - pass = false; - } - - // Check non-padding prefix matches expected exactly - for (size_t j = 0; j < exp.size() && j < tokens.size(); ++j) { - if (tokens[j] != exp[j]) { - fprintf(stderr, "FAIL [%d]: token[%zu] = %d, expected %d\n", - i, j, tokens[j], exp[j]); - pass = false; - } - } - - // Check remaining positions are zero-padded - for (size_t j = exp.size(); j < tokens.size(); ++j) { - if (tokens[j] != 0) { - fprintf(stderr, "FAIL [%d]: expected 0 at position %zu, got %d\n", - i, j, tokens[j]); - pass = false; - break; - } - } - - // Print result - fprintf(stderr, "%s [%d] \"%s\" → [", pass ? "PASS" : "FAIL", i, CASES[i].text); - for (size_t j = 0; j < tokens.size(); ++j) { - if (j > 0) fprintf(stderr, ", "); - fprintf(stderr, "%d", tokens[j]); - if (tokens[j] == 0 && j > 1) { fprintf(stderr, ", ..."); break; } - } - fprintf(stderr, "]\n"); - - if (pass) n_pass++; - } - - fprintf(stderr, "\n%d/%d tests passed\n", n_pass, N_CASES); - return (n_pass == N_CASES) ? 0 : 1; -} diff --git a/tests/test_utils.h b/tests/test_utils.h deleted file mode 100644 index ef5907a..0000000 --- a/tests/test_utils.h +++ /dev/null @@ -1,194 +0,0 @@ -#pragma once - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -struct ref_tensor_f32 { - std::vector data; - std::vector shape; - - int numel() const { - int n = 1; - for (int d : shape) { - n *= d; - } - return n; - } -}; - -struct ref_tensor_i32 { - std::vector data; - std::vector shape; - - int numel() const { - int n = 1; - for (int d : shape) { - n *= d; - } - return n; - } -}; - -struct compare_result { - float max_diff = 0.0f; - float mean_diff = 0.0f; - float cosine_sim = 0.0f; - int n_bad = 0; - int n_total = 0; - int worst_index = -1; - float worst_a = 0.0f; - float worst_b = 0.0f; -}; - -static inline std::vector load_shape_file(const std::string & path) { - std::vector shape; - std::ifstream f(path); - if (!f) { - return shape; - } - - std::string line; - std::getline(f, line); - size_t pos = 0; - while (pos < line.size()) { - size_t end = line.find(',', pos); - if (end == std::string::npos) { - end = line.size(); - } - shape.push_back(std::stoi(line.substr(pos, end - pos))); - pos = end + 1; - } - return shape; -} - -static inline ref_tensor_f32 load_ref_f32(const std::string & path) { - ref_tensor_f32 t; - t.shape = load_shape_file(path + ".shape"); - if (t.shape.empty()) { - return t; - } - - std::ifstream f(path + ".bin", std::ios::binary); - if (!f) { - t.shape.clear(); - return t; - } - - t.data.resize(t.numel()); - f.read(reinterpret_cast(t.data.data()), t.numel() * sizeof(float)); - return t; -} - -static inline ref_tensor_i32 load_ref_i32(const std::string & path) { - ref_tensor_i32 t; - t.shape = load_shape_file(path + ".shape"); - if (t.shape.empty()) { - return t; - } - - std::ifstream f(path + ".bin", std::ios::binary); - if (!f) { - t.shape.clear(); - return t; - } - - t.data.resize(t.numel()); - f.read(reinterpret_cast(t.data.data()), t.numel() * sizeof(int32_t)); - return t; -} - -static inline compare_result compare_tensors(const float * a, - const float * b, - int n, - float atol = 1e-4f) { - compare_result r; - r.n_total = n; - - double sum_diff = 0.0; - double dot_ab = 0.0; - double dot_aa = 0.0; - double dot_bb = 0.0; - - for (int i = 0; i < n; ++i) { - float diff = 0.0f; - if (!std::isfinite(a[i]) || !std::isfinite(b[i])) { - diff = (std::isinf(a[i]) && std::isinf(b[i]) && - ((a[i] > 0.0f) == (b[i] > 0.0f))) ? 0.0f : INFINITY; - } else { - diff = fabsf(a[i] - b[i]); - dot_ab += (double) a[i] * (double) b[i]; - dot_aa += (double) a[i] * (double) a[i]; - dot_bb += (double) b[i] * (double) b[i]; - } - - if (diff > r.max_diff) { - r.max_diff = diff; - r.worst_index = i; - r.worst_a = a[i]; - r.worst_b = b[i]; - } - if (diff > atol) { - r.n_bad++; - } - sum_diff += diff; - } - - r.mean_diff = n > 0 ? (float) (sum_diff / n) : 0.0f; - const double denom = sqrt(dot_aa) * sqrt(dot_bb); - r.cosine_sim = denom > 0.0 ? (float) (dot_ab / denom) : 0.0f; - return r; -} - -static inline int compare_exact_i32(const std::vector & got, - const std::vector & ref) { - const size_t n = got.size() < ref.size() ? got.size() : ref.size(); - int n_bad = got.size() == ref.size() ? 0 : 1; - for (size_t i = 0; i < n; ++i) { - if (got[i] != ref[i]) { - ++n_bad; - } - } - return n_bad; -} - -static inline bool ensure_dir(const std::string & path) { - if (path.empty()) { - return true; - } - - std::string cur; - size_t pos = 0; - if (path[0] == '/') { - cur = "/"; - pos = 1; - } - - while (pos <= path.size()) { - size_t next = path.find('/', pos); - if (next == std::string::npos) { - next = path.size(); - } - - const std::string part = path.substr(pos, next - pos); - if (!part.empty()) { - if (!cur.empty() && cur.back() != '/') { - cur.push_back('/'); - } - cur += part; - if (mkdir(cur.c_str(), 0755) != 0 && errno != EEXIST) { - return false; - } - } - - pos = next + 1; - } - - return true; -} diff --git a/tests/test_v_cont_removal.cpp b/tests/test_v_cont_removal.cpp deleted file mode 100644 index db185c9..0000000 --- a/tests/test_v_cont_removal.cpp +++ /dev/null @@ -1,183 +0,0 @@ -// Validate that removing cont(permute) for V produces identical output. -// Builds the same ViT block twice — once with V cont, once without — on CPU. -// Output must be bit-identical. - -#include "ggml-alloc.h" -#include "ggml-backend.h" -#include "ggml-cpu.h" -#include "ggml.h" - -#ifdef GGML_USE_METAL -#include "ggml-metal.h" -#endif - -#include -#include -#include - -struct block_result { - std::vector output; - int n_nodes; -}; - -static block_result run_vit_block(ggml_backend_t backend, bool use_v_cont) { - const int E = 1024, HD = 64, NH = 16, MLP = 4736; - const int W = 24, H = 24, NP = 9, N = W * H; - const ggml_type WT = GGML_TYPE_F16; - - size_t ctx_size = ggml_tensor_overhead() * 128 + ggml_graph_overhead_custom(128, false); - ggml_init_params params = {ctx_size, nullptr, true}; - auto* ctx = ggml_init(params); - std::vector inputs; - - auto inp = [&](ggml_type type, int64_t d0, int64_t d1=1, int64_t d2=1, int64_t d3=1) { - auto* t = ggml_new_tensor_4d(ctx, type, d0, d1, d2, d3); - ggml_set_input(t); - inputs.push_back(t); - return t; - }; - - auto* x = inp(GGML_TYPE_F32, E, W, H, NP); - auto* norm1_w = inp(GGML_TYPE_F32, E); - auto* norm1_b = inp(GGML_TYPE_F32, E); - auto* xn = ggml_add(ctx, ggml_mul(ctx, ggml_norm(ctx, x, 1e-6f), norm1_w), norm1_b); - - auto* qkv_w = inp(WT, E, 3*E); - auto* qkv_b = inp(GGML_TYPE_F32, 3*E); - auto* qkv = ggml_add(ctx, ggml_mul_mat(ctx, qkv_w, xn), qkv_b); - - qkv = ggml_reshape_4d(ctx, qkv, E, 3, N, NP); - qkv = ggml_cont(ctx, ggml_permute(ctx, qkv, 0, 3, 1, 2)); - auto* Q = ggml_view_3d(ctx, qkv, E, N, NP, qkv->nb[1], qkv->nb[2], 0); - auto* K = ggml_view_3d(ctx, qkv, E, N, NP, qkv->nb[1], qkv->nb[2], qkv->nb[3]); - auto* V = ggml_view_3d(ctx, qkv, E, N, NP, qkv->nb[1], qkv->nb[2], 2*qkv->nb[3]); - - // Q and K: always use cont (they need it for RoPE) - Q = ggml_reshape_4d(ctx, Q, HD, NH, N, NP); - Q = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 1, 3)); - Q = ggml_reshape_4d(ctx, Q, HD, N, NH, NP); - - K = ggml_reshape_4d(ctx, K, HD, NH, N, NP); - K = ggml_cont(ctx, ggml_permute(ctx, K, 0, 2, 1, 3)); - K = ggml_reshape_4d(ctx, K, HD, N, NH, NP); - - // V: the part we're testing - if (use_v_cont) { - // OLD path: reshape → permute → cont → reshape_3d → reshape_4d - V = ggml_reshape_4d(ctx, V, HD, NH, N, NP); - V = ggml_cont(ctx, ggml_permute(ctx, V, 0, 2, 1, 3)); - V = ggml_reshape_3d(ctx, V, HD, N, NH * NP); - V = ggml_reshape_4d(ctx, V, HD, N, NH, NP); - } else { - // NEW path: reshape → permute (no cont, no reshape) - V = ggml_reshape_4d(ctx, V, HD, NH, N, NP); - V = ggml_permute(ctx, V, 0, 2, 1, 3); - } - - auto* attn = ggml_flash_attn_ext(ctx, Q, K, V, nullptr, 1.0f/8.0f, 0, 0); - auto* attn_r = ggml_reshape_4d(ctx, attn, E, W, H, NP); - - auto* proj_w = inp(WT, E, E); - auto* proj_b = inp(GGML_TYPE_F32, E); - auto* proj = ggml_add(ctx, ggml_mul_mat(ctx, proj_w, attn_r), proj_b); - auto* res1 = ggml_add(ctx, x, proj); - - auto* norm2_w = inp(GGML_TYPE_F32, E); - auto* norm2_b = inp(GGML_TYPE_F32, E); - auto* xn2 = ggml_add(ctx, ggml_mul(ctx, ggml_norm(ctx, res1, 1e-6f), norm2_w), norm2_b); - - auto* fc1_w = inp(WT, E, MLP); - auto* fc1_b = inp(GGML_TYPE_F32, MLP); - auto* fc2_w = inp(WT, MLP, E); - auto* fc2_b = inp(GGML_TYPE_F32, E); - auto* h = ggml_add(ctx, ggml_mul_mat(ctx, fc1_w, xn2), fc1_b); - h = ggml_gelu_erf(ctx, h); - h = ggml_add(ctx, ggml_mul_mat(ctx, fc2_w, h), fc2_b); - auto* out = ggml_add(ctx, res1, h); - ggml_set_output(out); - - auto* graph = ggml_new_graph_custom(ctx, 128, false); - ggml_build_forward_expand(graph, out); - - auto* galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - ggml_gallocr_reserve(galloc, graph); - ggml_gallocr_alloc_graph(galloc, graph); - - // Fill with deterministic data - for (auto* t : inputs) { - int64_t nel = ggml_nelements(t); - if (t->type == GGML_TYPE_F32) { - std::vector d(nel); - for (int64_t j = 0; j < nel; j++) d[j] = 0.01f * ((j % 97) - 48); - ggml_backend_tensor_set(t, d.data(), 0, nel * sizeof(float)); - } else { - std::vector d(nel); - for (int64_t j = 0; j < nel; j++) d[j] = ggml_fp32_to_fp16(0.01f * ((j % 97) - 48)); - ggml_backend_tensor_set(t, d.data(), 0, nel * sizeof(uint16_t)); - } - } - - if (ggml_backend_is_cpu(backend)) - ggml_backend_cpu_set_n_threads(backend, 8); - - ggml_backend_graph_compute(backend, graph); - - int64_t nel = ggml_nelements(out); - block_result r; - r.output.resize(nel); - r.n_nodes = ggml_graph_n_nodes(graph); - ggml_backend_tensor_get(out, r.output.data(), 0, nel * sizeof(float)); - - ggml_gallocr_free(galloc); - ggml_free(ctx); - return r; -} - -static bool compare(const std::vector& a, const std::vector& b, - const char* label, float tol) { - float max_diff = 0; - int max_idx = 0; - for (size_t i = 0; i < a.size(); i++) { - float d = std::fabs(a[i] - b[i]); - if (d > max_diff) { max_diff = d; max_idx = (int)i; } - } - bool ok = max_diff <= tol; - fprintf(stderr, " %s [%s]: max_diff=%.8f at [%d] (a=%.6f b=%.6f) (%zu elements)\n", - ok ? "PASS" : "FAIL", label, max_diff, max_idx, - a[max_idx], b[max_idx], a.size()); - return ok; -} - -int main() { - auto* cpu = ggml_backend_cpu_init(); - bool all_pass = true; - - // TEST 1: CPU with V cont vs CPU without V cont — must be IDENTICAL - fprintf(stderr, "\n=== CPU: with V cont vs without V cont ===\n"); - auto cpu_with = run_vit_block(cpu, true); - auto cpu_without = run_vit_block(cpu, false); - fprintf(stderr, " Nodes: with=%d, without=%d\n", cpu_with.n_nodes, cpu_without.n_nodes); - all_pass &= compare(cpu_with.output, cpu_without.output, "CPU with_cont vs without_cont", 0.0f); - -#ifdef GGML_USE_METAL - auto* metal = ggml_backend_metal_init(); - if (metal) { - // TEST 2: Metal with V cont vs Metal without V cont — must be IDENTICAL - fprintf(stderr, "\n=== Metal: with V cont vs without V cont ===\n"); - auto metal_with = run_vit_block(metal, true); - auto metal_without = run_vit_block(metal, false); - fprintf(stderr, " Nodes: with=%d, without=%d\n", metal_with.n_nodes, metal_without.n_nodes); - all_pass &= compare(metal_with.output, metal_without.output, "Metal with_cont vs without_cont", 0.0f); - - // TEST 3: CPU without vs Metal without — F16 precision diff (informational) - fprintf(stderr, "\n=== CPU vs Metal (both without V cont) ===\n"); - compare(cpu_without.output, metal_without.output, "CPU vs Metal (informational)", 50.0f); - - ggml_backend_free(metal); - } -#endif - - ggml_backend_free(cpu); - fprintf(stderr, "\n%s\n", all_pass ? "ALL TESTS PASSED" : "SOME TESTS FAILED"); - return all_pass ? 0 : 1; -} diff --git a/tests/test_video_2frame.cpp b/tests/test_video_2frame.cpp deleted file mode 100644 index 31e397c..0000000 --- a/tests/test_video_2frame.cpp +++ /dev/null @@ -1,107 +0,0 @@ -/** - * Simple 2-frame video tracking test. - * Frame 0: encode image, add instance via point prompt - * Frame 1: track (propagate mask to same image) - * - * Usage: test_video_2frame [x y] - */ -#include "sam3.h" -#include -#include - -int main(int argc, char ** argv) { - if (argc < 3) { - fprintf(stderr, "Usage: %s [x y]\n", argv[0]); - return 1; - } - - const std::string model_path = argv[1]; - const std::string image_path = argv[2]; - float px = argc > 4 ? atof(argv[3]) : 315.0f; - float py = argc > 4 ? atof(argv[4]) : 250.0f; - - sam3_params params; - params.model_path = model_path; - params.use_gpu = false; - params.n_threads = 4; - - fprintf(stderr, "Loading model...\n"); - auto model = sam3_load_model(params); - if (!model) { fprintf(stderr, "Failed to load model\n"); return 1; } - - // Create tracker (no text prompt → propagation only, no PCS detection) - sam3_video_params vp; - vp.hotstart_delay = 0; // instant confirmation (no warmup) - vp.max_keep_alive = 100; - auto tracker = sam3_create_tracker(*model, vp); - if (!tracker) { fprintf(stderr, "Failed to create tracker\n"); return 1; } - - auto state = sam3_create_state(*model, params); - if (!state) { fprintf(stderr, "Failed to create state\n"); return 1; } - - // Load image (use same image for both frames for simplicity) - auto image = sam3_load_image(image_path); - if (image.data.empty()) { fprintf(stderr, "Failed to load image\n"); return 1; } - fprintf(stderr, "Image: %dx%d\n", image.width, image.height); - - // ════════════════════════════════════════════════════════════════ - // Frame 0: Encode + add instance via point - // ════════════════════════════════════════════════════════════════ - fprintf(stderr, "\n═══ Frame 0: Encode + Add Instance ═══\n"); - - // Encode image - if (!sam3_encode_image(*state, *model, image)) { - fprintf(stderr, "Failed to encode frame 0\n"); return 1; - } - - // Add instance at the clicked point - sam3_pvs_params pvs; - pvs.pos_points.push_back({px, py}); - pvs.multimask = false; - - int inst_id = sam3_tracker_add_instance(*tracker, *state, *model, pvs); - fprintf(stderr, "Added instance %d at (%.1f, %.1f)\n", inst_id, px, py); - - if (inst_id < 0) { - fprintf(stderr, "Failed to add instance\n"); - return 1; - } - - // Check tracker state - fprintf(stderr, "Tracker frame: %d\n", sam3_tracker_frame_index(*tracker)); - - // ════════════════════════════════════════════════════════════════ - // Frame 1: Track (propagate to same image) - // ════════════════════════════════════════════════════════════════ - fprintf(stderr, "\n═══ Frame 1: Track ═══\n"); - - // Re-encode same image as "frame 1" - if (!sam3_encode_image(*state, *model, image)) { - fprintf(stderr, "Failed to encode frame 1\n"); return 1; - } - - // Track - auto result = sam3_track_frame(*tracker, *state, *model, image); - fprintf(stderr, "Track result: %zu detections\n", result.detections.size()); - - for (size_t i = 0; i < result.detections.size(); ++i) { - const auto& d = result.detections[i]; - fprintf(stderr, " det %zu: inst=%d score=%.4f iou=%.4f box=[%.1f,%.1f,%.1f,%.1f]\n", - i, d.instance_id, d.score, d.iou_score, - d.box.x0, d.box.y0, d.box.x1, d.box.y1); - - // Count mask pixels - int n_on = 0; - for (size_t j = 0; j < d.mask.data.size(); ++j) - if (d.mask.data[j] > 0) n_on++; - fprintf(stderr, " mask: %d/%zu pixels (%.1f%%)\n", - n_on, d.mask.data.size(), 100.0f * n_on / d.mask.data.size()); - - std::string mask_path = "/tmp/video_frame1_mask_" + std::to_string(i) + ".png"; - sam3_save_mask(d.mask, mask_path); - fprintf(stderr, " saved: %s\n", mask_path.c_str()); - } - - sam3_free_model(*model); - return result.detections.empty() ? 1 : 0; -} diff --git a/tests/test_video_text.cpp b/tests/test_video_text.cpp deleted file mode 100644 index 6829f2c..0000000 --- a/tests/test_video_text.cpp +++ /dev/null @@ -1,89 +0,0 @@ -/** - * Video tracking test with text prompt on first N frames. - * Usage: test_video_text [prompt] [n_frames] - */ -#include "sam3.h" -#include -#include - -int main(int argc, char ** argv) { - if (argc < 3) { - fprintf(stderr, "Usage: %s [prompt] [n_frames]\n", argv[0]); - return 1; - } - - const std::string model_path = argv[1]; - const std::string video_path = argv[2]; - const std::string prompt = argc > 3 ? argv[3] : "car"; - const int n_frames = argc > 4 ? atoi(argv[4]) : 5; - - sam3_params params; - params.model_path = model_path; - params.use_gpu = false; - params.n_threads = 4; - - fprintf(stderr, "Loading model...\n"); - auto model = sam3_load_model(params); - if (!model) { fprintf(stderr, "Failed to load model\n"); return 1; } - - auto state = sam3_create_state(*model, params); - if (!state) { fprintf(stderr, "Failed to create state\n"); return 1; } - - auto vi = sam3_get_video_info(video_path); - fprintf(stderr, "Video: %dx%d, %d frames, %.1f fps\n", - vi.width, vi.height, vi.n_frames, vi.fps); - - // Create tracker with text prompt - sam3_video_params vp; - vp.text_prompt = prompt; - vp.score_threshold = 0.3f; - vp.hotstart_delay = 3; // Quick confirmation for testing - auto tracker = sam3_create_tracker(*model, vp); - if (!tracker) { fprintf(stderr, "Failed to create tracker\n"); return 1; } - - int total_tracked = 0; - for (int fi = 0; fi < std::min(n_frames, vi.n_frames); ++fi) { - fprintf(stderr, "\n════════════════════════════════════════\n"); - fprintf(stderr, " Frame %d/%d\n", fi, n_frames); - fprintf(stderr, "════════════════════════════════════════\n"); - - auto frame = sam3_decode_video_frame(video_path, fi); - if (frame.data.empty()) { - fprintf(stderr, "Failed to decode frame %d\n", fi); - continue; - } - - auto result = sam3_track_frame(*tracker, *state, *model, frame); - fprintf(stderr, " Result: %zu detections\n", result.detections.size()); - - for (size_t i = 0; i < result.detections.size(); ++i) { - const auto& d = result.detections[i]; - int n_on = 0; - for (size_t j = 0; j < d.mask.data.size(); ++j) - if (d.mask.data[j] > 0) n_on++; - fprintf(stderr, " det %zu: inst=%d score=%.3f box=[%.0f,%.0f,%.0f,%.0f] mask=%.1f%%\n", - i, d.instance_id, d.score, - d.box.x0, d.box.y0, d.box.x1, d.box.y1, - 100.0f * n_on / d.mask.data.size()); - } - - if (!result.detections.empty()) { - total_tracked++; - // Save first frame's masks - if (fi == 0) { - for (size_t i = 0; i < result.detections.size() && i < 3; ++i) { - std::string path = "/tmp/video_text_f0_mask" + std::to_string(i) + ".png"; - sam3_save_mask(result.detections[i].mask, path); - fprintf(stderr, " saved: %s\n", path.c_str()); - } - } - } - } - - fprintf(stderr, "\n═══ Summary ═══\n"); - fprintf(stderr, " Prompt: '%s'\n", prompt.c_str()); - fprintf(stderr, " Frames with detections: %d/%d\n", total_tracked, n_frames); - - sam3_free_model(*model); - return total_tracked > 0 ? 0 : 1; -} diff --git a/tests/test_visual_only.cpp b/tests/test_visual_only.cpp deleted file mode 100644 index 584b59a..0000000 --- a/tests/test_visual_only.cpp +++ /dev/null @@ -1,102 +0,0 @@ -#include "sam3.h" -#include - -int main(int argc, char** argv) { - if (argc < 2) { - fprintf(stderr, "Usage: %s [image.jpg]\n", argv[0]); - return 1; - } - - sam3_params params; - params.model_path = argv[1]; - params.n_threads = 4; - - // 1. Load visual-only model - fprintf(stderr, "=== Test 1: Load visual-only model ===\n"); - auto model = sam3_load_model(params); - if (!model) { - fprintf(stderr, "FAIL: failed to load model\n"); - return 1; - } - fprintf(stderr, "PASS: model loaded\n"); - - // 2. Verify visual_only flag - fprintf(stderr, "\n=== Test 2: sam3_is_visual_only ===\n"); - if (!sam3_is_visual_only(*model)) { - fprintf(stderr, "FAIL: expected visual_only=true\n"); - return 1; - } - fprintf(stderr, "PASS: sam3_is_visual_only() = true\n"); - - // 3. PCS should return empty result (no crash) - fprintf(stderr, "\n=== Test 3: PCS guard ===\n"); - auto state = sam3_create_state(*model, params); - if (!state) { - fprintf(stderr, "FAIL: create_state\n"); - return 1; - } - sam3_pcs_params pcs; - pcs.text_prompt = "cat"; - auto r = sam3_segment_pcs(*state, *model, pcs); - if (!r.detections.empty()) { - fprintf(stderr, "FAIL: PCS should return empty on visual-only\n"); - return 1; - } - fprintf(stderr, "PASS: PCS returns empty on visual-only model\n"); - - // 4. If image provided, test PVS + visual tracking - if (argc >= 3) { - fprintf(stderr, "\n=== Test 4: PVS on visual-only model ===\n"); - auto img = sam3_load_image(argv[2]); - if (img.data.empty()) { - fprintf(stderr, "FAIL: failed to load image '%s'\n", argv[2]); - return 1; - } - if (!sam3_encode_image(*state, *model, img)) { - fprintf(stderr, "FAIL: encode_image\n"); - return 1; - } - - // PVS with center point - sam3_pvs_params pvs; - pvs.pos_points.push_back({(float)img.width / 2, (float)img.height / 2}); - auto rp = sam3_segment_pvs(*state, *model, pvs); - if (rp.detections.empty()) { - fprintf(stderr, "FAIL: PVS returned no detections\n"); - return 1; - } - fprintf(stderr, "PASS: PVS produced %zu detection(s)\n", rp.detections.size()); - - // 5. Visual tracker: create, add instance, propagate - fprintf(stderr, "\n=== Test 5: Visual tracker ===\n"); - sam3_visual_track_params vtp; - auto tracker = sam3_create_visual_tracker(*model, vtp); - if (!tracker) { - fprintf(stderr, "FAIL: create_visual_tracker\n"); - return 1; - } - - int id = sam3_tracker_add_instance(*tracker, *state, *model, pvs); - if (id < 0) { - fprintf(stderr, "FAIL: tracker_add_instance returned %d\n", id); - return 1; - } - fprintf(stderr, "PASS: tracker_add_instance -> id=%d\n", id); - - // Propagate same frame as a smoke test - auto rr = sam3_propagate_frame(*tracker, *state, *model, img); - fprintf(stderr, "PASS: propagate_frame -> %zu detection(s)\n", rr.detections.size()); - - // 6. track_frame should fail on visual-only - fprintf(stderr, "\n=== Test 6: track_frame guard ===\n"); - auto rt = sam3_track_frame(*tracker, *state, *model, img); - if (!rt.detections.empty()) { - fprintf(stderr, "FAIL: track_frame should return empty on visual-only\n"); - return 1; - } - fprintf(stderr, "PASS: track_frame returns empty on visual-only model\n"); - } - - fprintf(stderr, "\nAll tests passed.\n"); - return 0; -} diff --git a/tests/test_visual_only_compare.cpp b/tests/test_visual_only_compare.cpp deleted file mode 100644 index 7d12c00..0000000 --- a/tests/test_visual_only_compare.cpp +++ /dev/null @@ -1,260 +0,0 @@ -/** - * Compare visual-only model output vs full model output for PVS (points/boxes) - * and video propagation. The tracker-path weights are identical, so outputs - * should match bit-for-bit (same ggml graph, same weights, same inputs). - * - * Usage: ./test_visual_only_compare [video.mp4] - */ -#include "sam3.h" - -#include -#include -#include -#include -#include -#include - -static float mask_iou(const sam3_mask& a, const sam3_mask& b) { - if (a.width != b.width || a.height != b.height) return 0.0f; - int inter = 0, uni = 0; - for (int i = 0; i < (int)a.data.size(); ++i) { - bool fa = a.data[i] > 127; - bool fb = b.data[i] > 127; - if (fa && fb) inter++; - if (fa || fb) uni++; - } - return uni > 0 ? (float)inter / uni : 1.0f; -} - -static int count_fg(const sam3_mask& m) { - int n = 0; - for (auto v : m.data) if (v > 127) n++; - return n; -} - -static bool compare_results(const char* label, - const sam3_result& full_r, - const sam3_result& vis_r) { - fprintf(stderr, "\n--- %s ---\n", label); - if (full_r.detections.size() != vis_r.detections.size()) { - fprintf(stderr, " FAIL: detection count mismatch: full=%zu visual=%zu\n", - full_r.detections.size(), vis_r.detections.size()); - return false; - } - if (full_r.detections.empty()) { - fprintf(stderr, " WARN: both returned 0 detections\n"); - return true; - } - - bool ok = true; - for (size_t i = 0; i < full_r.detections.size(); ++i) { - const auto& fd = full_r.detections[i]; - const auto& vd = vis_r.detections[i]; - - float iou = mask_iou(fd.mask, vd.mask); - int fg_full = count_fg(fd.mask); - int fg_vis = count_fg(vd.mask); - float score_diff = fabsf(fd.score - vd.score); - float iou_diff = fabsf(fd.iou_score - vd.iou_score); - - fprintf(stderr, " det[%zu]: mask_iou=%.6f fg_full=%d fg_vis=%d " - "score_diff=%.6f iou_diff=%.6f\n", - i, iou, fg_full, fg_vis, score_diff, iou_diff); - - if (iou < 0.999f) { - fprintf(stderr, " FAIL: mask IoU %.6f < 0.999 threshold\n", iou); - ok = false; - } - if (score_diff > 1e-4f) { - fprintf(stderr, " FAIL: score diff %.6f > 1e-4\n", score_diff); - ok = false; - } - } - fprintf(stderr, " %s\n", ok ? "PASS" : "FAIL"); - return ok; -} - -int main(int argc, char** argv) { - if (argc < 4) { - fprintf(stderr, "Usage: %s [video.mp4]\n", argv[0]); - return 1; - } - - const std::string full_path = argv[1]; - const std::string vis_path = argv[2]; - const std::string image_path = argv[3]; - const std::string video_path = (argc >= 5) ? argv[4] : ""; - - int n_pass = 0, n_fail = 0; - auto tally = [&](bool ok) { if (ok) n_pass++; else n_fail++; }; - - auto img = sam3_load_image(image_path); - if (img.data.empty()) { - fprintf(stderr, "Failed to load image '%s'\n", image_path.c_str()); - return 1; - } - fprintf(stderr, "Image: %dx%d\n", img.width, img.height); - - // ── Define test prompts ───────────────────────────────────────────── - struct pvs_test { - const char* name; - sam3_pvs_params params; - }; - - pvs_test tests[] = { - {"center_point", {}}, - {"box_prompt", {}}, - {"two_points", {}}, - {"point_and_box", {}}, - }; - - // Center point - tests[0].params.pos_points.push_back({(float)img.width / 2, (float)img.height / 2}); - - // Box prompt - tests[1].params.box = {(float)img.width * 0.1f, (float)img.height * 0.1f, - (float)img.width * 0.9f, (float)img.height * 0.9f}; - tests[1].params.use_box = true; - - // Two points (positive + negative) - tests[2].params.pos_points.push_back({(float)img.width / 2, (float)img.height / 2}); - tests[2].params.neg_points.push_back({10.0f, 10.0f}); - - // Point + box - tests[3].params.pos_points.push_back({(float)img.width / 2, (float)img.height / 2}); - tests[3].params.box = {(float)img.width * 0.2f, (float)img.height * 0.2f, - (float)img.width * 0.8f, (float)img.height * 0.8f}; - tests[3].params.use_box = true; - - // ═══════════════════════════════════════════════════════════════════════ - // Run on full model - // ═══════════════════════════════════════════════════════════════════════ - fprintf(stderr, "\n═══ Loading FULL model ═══\n"); - sam3_params fparams; - fparams.model_path = full_path; - fparams.n_threads = 4; - fparams.use_gpu = false; - - auto full_model = sam3_load_model(fparams); - if (!full_model) { fprintf(stderr, "Failed to load full model\n"); return 1; } - assert(!sam3_is_visual_only(*full_model)); - - auto full_state = sam3_create_state(*full_model, fparams); - if (!full_state) { fprintf(stderr, "Failed to create state\n"); return 1; } - - if (!sam3_encode_image(*full_state, *full_model, img)) { - fprintf(stderr, "Failed to encode image (full)\n"); return 1; - } - - std::vector full_results; - for (auto& t : tests) { - fprintf(stderr, " Running PVS [%s] on full model...\n", t.name); - full_results.push_back(sam3_segment_pvs(*full_state, *full_model, t.params)); - } - - // Video propagation on full model - std::vector full_video_results; - if (!video_path.empty()) { - fprintf(stderr, "\n Running video propagation on full model...\n"); - sam3_visual_track_params vtp; - auto tracker = sam3_create_visual_tracker(*full_model, vtp); - - // Frame 0: encode + add instance - auto f0 = sam3_decode_video_frame(video_path, 0); - if (!f0.data.empty()) { - sam3_encode_image(*full_state, *full_model, f0); - sam3_pvs_params vpvs; - vpvs.pos_points.push_back({(float)f0.width / 2, (float)f0.height / 2}); - sam3_tracker_add_instance(*tracker, *full_state, *full_model, vpvs); - - // Propagate frames 1-3 - for (int fi = 1; fi <= 3; ++fi) { - auto frame = sam3_decode_video_frame(video_path, fi); - if (frame.data.empty()) break; - full_video_results.push_back( - sam3_propagate_frame(*tracker, *full_state, *full_model, frame)); - } - } - } - - // Free full model to save memory - full_state.reset(); - sam3_free_model(*full_model); - full_model.reset(); - - // ═══════════════════════════════════════════════════════════════════════ - // Run on visual-only model - // ═══════════════════════════════════════════════════════════════════════ - fprintf(stderr, "\n═══ Loading VISUAL-ONLY model ═══\n"); - sam3_params vparams; - vparams.model_path = vis_path; - vparams.n_threads = 4; - vparams.use_gpu = false; - - auto vis_model = sam3_load_model(vparams); - if (!vis_model) { fprintf(stderr, "Failed to load visual model\n"); return 1; } - assert(sam3_is_visual_only(*vis_model)); - - auto vis_state = sam3_create_state(*vis_model, vparams); - if (!vis_state) { fprintf(stderr, "Failed to create state\n"); return 1; } - - if (!sam3_encode_image(*vis_state, *vis_model, img)) { - fprintf(stderr, "Failed to encode image (visual)\n"); return 1; - } - - std::vector vis_results; - for (auto& t : tests) { - fprintf(stderr, " Running PVS [%s] on visual-only model...\n", t.name); - vis_results.push_back(sam3_segment_pvs(*vis_state, *vis_model, t.params)); - } - - // Video propagation on visual-only model - std::vector vis_video_results; - if (!video_path.empty()) { - fprintf(stderr, "\n Running video propagation on visual-only model...\n"); - sam3_visual_track_params vtp; - auto tracker = sam3_create_visual_tracker(*vis_model, vtp); - - auto f0 = sam3_decode_video_frame(video_path, 0); - if (!f0.data.empty()) { - sam3_encode_image(*vis_state, *vis_model, f0); - sam3_pvs_params vpvs; - vpvs.pos_points.push_back({(float)f0.width / 2, (float)f0.height / 2}); - sam3_tracker_add_instance(*tracker, *vis_state, *vis_model, vpvs); - - for (int fi = 1; fi <= 3; ++fi) { - auto frame = sam3_decode_video_frame(video_path, fi); - if (frame.data.empty()) break; - vis_video_results.push_back( - sam3_propagate_frame(*tracker, *vis_state, *vis_model, frame)); - } - } - } - - vis_state.reset(); - sam3_free_model(*vis_model); - vis_model.reset(); - - // ═══════════════════════════════════════════════════════════════════════ - // Compare results - // ═══════════════════════════════════════════════════════════════════════ - fprintf(stderr, "\n══════════════════════════════════════════════════\n"); - fprintf(stderr, " COMPARISON: Full model vs Visual-only model\n"); - fprintf(stderr, "══════════════════════════════════════════════════\n"); - - for (size_t i = 0; i < full_results.size(); ++i) { - tally(compare_results(tests[i].name, full_results[i], vis_results[i])); - } - - for (size_t i = 0; i < full_video_results.size() && i < vis_video_results.size(); ++i) { - char label[64]; - snprintf(label, sizeof(label), "video_frame_%zu", i + 1); - tally(compare_results(label, full_video_results[i], vis_video_results[i])); - } - - fprintf(stderr, "\n══════════════════════════════════════════════════\n"); - fprintf(stderr, " TOTAL: %d PASS, %d FAIL\n", n_pass, n_fail); - fprintf(stderr, "══════════════════════════════════════════════════\n\n"); - - return n_fail > 0 ? 1 : 0; -} diff --git a/tests/test_vit.cpp b/tests/test_vit.cpp deleted file mode 100644 index affb76f..0000000 --- a/tests/test_vit.cpp +++ /dev/null @@ -1,465 +0,0 @@ -#include "sam3.h" - -#include -#include -#include -#include -#include -#include -#include - -// ═══════════════════════════════════════════════════════════════════════════════ -// Helpers: load reference tensors -// ═══════════════════════════════════════════════════════════════════════════════ - -struct ref_tensor { - std::vector data; - std::vector shape; - int numel() const { - int n = 1; - for (auto d : shape) n *= d; - return n; - } -}; - -static ref_tensor load_ref(const std::string & path) { - ref_tensor t; - // Load shape - { - std::ifstream f(path + ".shape"); - if (!f) { - fprintf(stderr, " [SKIP] %s.shape not found\n", path.c_str()); - return t; - } - std::string line; - std::getline(f, line); - size_t pos = 0; - while (pos < line.size()) { - size_t end = line.find(',', pos); - if (end == std::string::npos) end = line.size(); - t.shape.push_back(std::stoi(line.substr(pos, end - pos))); - pos = end + 1; - } - } - // Load data - { - std::ifstream f(path + ".bin", std::ios::binary); - if (!f) { - fprintf(stderr, " [SKIP] %s.bin not found\n", path.c_str()); - t.shape.clear(); - return t; - } - t.data.resize(t.numel()); - f.read(reinterpret_cast(t.data.data()), t.numel() * sizeof(float)); - } - return t; -} - -// Compare two float arrays with tolerance. -// Returns max absolute difference. -static float compare_tensors(const float * a, const float * b, int n, - float atol = 1e-4f, float rtol = 1e-3f) { - float max_diff = 0.0f; - int n_bad = 0; - for (int i = 0; i < n; ++i) { - float diff = fabsf(a[i] - b[i]); - float threshold = atol + rtol * fabsf(b[i]); - if (diff > threshold) n_bad++; - if (diff > max_diff) max_diff = diff; - } - return max_diff; -} - -static bool check(const std::string & name, const float * got, const ref_tensor & ref, - float atol = 1e-4f, float rtol = 1e-3f) { - if (ref.data.empty()) { - fprintf(stderr, " [SKIP] %s — no reference data\n", name.c_str()); - return true; - } - float max_diff = compare_tensors(got, ref.data.data(), ref.numel(), atol, rtol); - bool ok = max_diff < atol + rtol; - fprintf(stderr, " %s %s: max_diff=%.6f (atol=%.1e rtol=%.1e)\n", - ok ? "[PASS]" : "[FAIL]", name.c_str(), max_diff, atol, rtol); - return ok; -} - -// ═══════════════════════════════════════════════════════════════════════════════ -// Test: preprocessing -// ═══════════════════════════════════════════════════════════════════════════════ - -static bool test_preprocessing(const std::string & image_path, const std::string & ref_dir) { - fprintf(stderr, "\n=== Test: Preprocessing ===\n"); - - // Load image - auto img = sam3_load_image(image_path); - if (img.data.empty()) { - fprintf(stderr, " Failed to load image: %s\n", image_path.c_str()); - return false; - } - fprintf(stderr, " Loaded %dx%d image\n", img.width, img.height); - - // Load reference - auto ref = load_ref(ref_dir + "/preprocessed"); - if (ref.data.empty()) return true; // skip - - // The reference is [1, 3, 1008, 1008] from PyTorch (NCHW) - // We check a reasonable tolerance since bilinear resize implementations differ slightly - fprintf(stderr, " Reference shape: "); - for (auto d : ref.shape) fprintf(stderr, "%d ", d); - fprintf(stderr, "\n"); - - // Note: we can't directly test preprocessing since it's internal to sam3_encode_image. - // For now, just verify the image loaded correctly. - fprintf(stderr, " [INFO] Preprocessing test requires running encode_image — tested implicitly\n"); - return true; -} - -// ═══════════════════════════════════════════════════════════════════════════════ -// Test: RoPE computation -// ═══════════════════════════════════════════════════════════════════════════════ - -// Forward-declare the internal function (defined in sam3.cpp) -// We'll test it by replicating the computation here. -static bool test_rope(const std::string & ref_dir) { - fprintf(stderr, "\n=== Test: RoPE ===\n"); - - // Test window RoPE (24x24) - auto ref_window = load_ref(ref_dir + "/rope_window_real"); - if (!ref_window.data.empty()) { - // ref is [576, 32, 2] — (cos, sin) pairs - // Our compute_axial_cis produces [N, head_dim] with interleaved (cos, sin) - // But the Python ref stores them as [N, half_head, 2] where [:,:,0]=cos, [:,:,1]=sin - const int N = 576; - const int half_head = 32; - const int dim = 64; - const float theta = 10000.0f; - const float scale_pos = 1.0f; - - std::vector our_rope(N * dim); - - const int half_dim = dim / 4; // 16 - std::vector freqs(half_dim); - for (int i = 0; i < half_dim; ++i) { - freqs[i] = 1.0f / powf(theta, (float)(i * 4) / dim); - } - - for (int idx = 0; idx < N; ++idx) { - float t_x = (float)(idx % 24) * scale_pos; - float t_y = (float)(idx / 24) * scale_pos; - - for (int i = 0; i < half_dim; ++i) { - float angle_x = t_x * freqs[i]; - our_rope[idx * dim + i * 2 + 0] = cosf(angle_x); - our_rope[idx * dim + i * 2 + 1] = sinf(angle_x); - } - for (int i = 0; i < half_dim; ++i) { - float angle_y = t_y * freqs[i]; - our_rope[idx * dim + half_dim * 2 + i * 2 + 0] = cosf(angle_y); - our_rope[idx * dim + half_dim * 2 + i * 2 + 1] = sinf(angle_y); - } - } - - // Compare: ref is [N, 32, 2] = [N, half_head, (cos,sin)] - // Our layout is [N, dim] = [N, 64] with pairs (cos_x, sin_x, cos_x, sin_x, ..., cos_y, sin_y, ...) - // Need to rearrange comparison: - // ref[idx, j, 0] = cos of freq j for position idx - // ref[idx, j, 1] = sin of freq j for position idx - // Python: freqs_cis = cat([freqs_cis_x, freqs_cis_y], dim=-1) → [N, 32] - // freqs_cis_x has 16 complex values, freqs_cis_y has 16 complex values → total 32 - // view_as_real → [N, 32, 2] where [:,i,0]=cos, [:,i,1]=sin - // So ref[idx, j, 0] for j<16 = cos(t_x * freqs[j]) - // ref[idx, j, 1] for j<16 = sin(t_x * freqs[j]) - // ref[idx, j, 0] for j>=16 = cos(t_y * freqs[j-16]) - // ref[idx, j, 1] for j>=16 = sin(t_y * freqs[j-16]) - - // Our layout: [idx * 64 + j*2+0] = cos, [idx * 64 + j*2+1] = sin - // For x: j = 0..15 → our[idx*64 + j*2], ref[idx*64 + j*2] - // For y: j = 16..31 → our[idx*64 + 32 + (j-16)*2], ref[idx*64 + (j)*2] - - // Actually both are the same layout! ref is stored flat as [N*32*2] = [N*64] - // and our_rope is [N*64] with the same interleaving. - float max_diff = compare_tensors(our_rope.data(), ref_window.data.data(), N * dim); - bool ok = max_diff < 1e-5f; - fprintf(stderr, " %s window RoPE (24x24): max_diff=%.6e\n", ok ? "[PASS]" : "[FAIL]", max_diff); - } - - // Test global RoPE (72x72) - auto ref_global = load_ref(ref_dir + "/rope_global_real"); - if (!ref_global.data.empty()) { - const int N = 5184; - const int dim = 64; - const float theta = 10000.0f; - const float scale_pos = 24.0f / 72.0f; // 1/3 - - std::vector our_rope(N * dim); - const int half_dim = dim / 4; - std::vector freqs(half_dim); - for (int i = 0; i < half_dim; ++i) { - freqs[i] = 1.0f / powf(theta, (float)(i * 4) / dim); - } - - for (int idx = 0; idx < N; ++idx) { - float t_x = (float)(idx % 72) * scale_pos; - float t_y = (float)(idx / 72) * scale_pos; - - for (int i = 0; i < half_dim; ++i) { - float angle_x = t_x * freqs[i]; - our_rope[idx * dim + i * 2 + 0] = cosf(angle_x); - our_rope[idx * dim + i * 2 + 1] = sinf(angle_x); - } - for (int i = 0; i < half_dim; ++i) { - float angle_y = t_y * freqs[i]; - our_rope[idx * dim + half_dim * 2 + i * 2 + 0] = cosf(angle_y); - our_rope[idx * dim + half_dim * 2 + i * 2 + 1] = sinf(angle_y); - } - } - - float max_diff = compare_tensors(our_rope.data(), ref_global.data.data(), N * dim); - bool ok = max_diff < 1e-5f; - fprintf(stderr, " %s global RoPE (72x72): max_diff=%.6e\n", ok ? "[PASS]" : "[FAIL]", max_diff); - } - - return true; -} - -// ═══════════════════════════════════════════════════════════════════════════════ -// Test: sinusoidal PE -// ═══════════════════════════════════════════════════════════════════════════════ - -// Forward-declare for testing -extern "C" { - // We can't easily call static functions from sam3.cpp, so we'll replicate the logic -} - -static bool test_sinusoidal_pe(const std::string & ref_dir) { - fprintf(stderr, "\n=== Test: Sinusoidal PE ===\n"); - - struct { int H; int W; const char * name; } cases[] = { - {288, 288, "pe_288"}, - {144, 144, "pe_144"}, - { 72, 72, "pe_72"}, - { 36, 36, "pe_36"}, - }; - - for (auto & tc : cases) { - auto ref = load_ref(ref_dir + "/" + tc.name); - if (ref.data.empty()) continue; - - // Compute our sinusoidal PE - const int H = tc.H, W = tc.W, d_model = 256; - const int half = d_model / 2; - const float scale = 2.0f * (float)M_PI; - const float temperature = 10000.0f; - - // ref is [1, 256, H, W] in NCHW from PyTorch - std::vector our_pe(d_model * H * W); - - for (int y = 0; y < H; ++y) { - for (int x = 0; x < W; ++x) { - float pos_y = ((float)(y + 1) / (float)(H)) * scale; - float pos_x = ((float)(x + 1) / (float)(W)) * scale; - - for (int i = 0; i < half; ++i) { - float dim_t = powf(temperature, 2.0f * (float)(i / 2) / (float)half); - - float val_x, val_y; - if (i % 2 == 0) { - val_x = sinf(pos_x / dim_t); - val_y = sinf(pos_y / dim_t); - } else { - val_x = cosf(pos_x / dim_t); - val_y = cosf(pos_y / dim_t); - } - - // PyTorch output is [1, 256, H, W] = [N, C, H, W] - // Channel layout: first half is pos_y, second half is pos_x - // Index: [0, c, y, x] → c * H * W + y * W + x - our_pe[i * H * W + y * W + x] = val_y; - our_pe[(i + half) * H * W + y * W + x] = val_x; - } - } - } - - float max_diff = compare_tensors(our_pe.data(), ref.data.data(), d_model * H * W); - bool ok = max_diff < 1e-5f; - fprintf(stderr, " %s %s (%dx%d): max_diff=%.6e\n", - ok ? "[PASS]" : "[FAIL]", tc.name, tc.H, tc.W, max_diff); - } - - return true; -} - -// ═══════════════════════════════════════════════════════════════════════════════ -// Test: full image encoding -// ═══════════════════════════════════════════════════════════════════════════════ - -static bool test_encode_image(const std::string & model_path, - const std::string & image_path, - const std::string & ref_dir) { - fprintf(stderr, "\n=== Test: Full Image Encoding ===\n"); - - sam3_params params; - params.model_path = model_path; - params.use_gpu = false; // CPU for determinism - params.n_threads = 4; - - auto model = sam3_load_model(params); - if (!model) { - fprintf(stderr, " Failed to load model\n"); - return false; - } - - auto state = sam3_create_state(*model, params); - if (!state) { - fprintf(stderr, " Failed to create state\n"); - return false; - } - - auto img = sam3_load_image(image_path); - if (img.data.empty()) { - fprintf(stderr, " Failed to load image\n"); - return false; - } - - bool ok = sam3_encode_image(*state, *model, img); - if (!ok) { - fprintf(stderr, " sam3_encode_image failed\n"); - return false; - } - - fprintf(stderr, " Image encoding completed\n"); - - int n_fail_enc = 0; - std::string dump_dir = ref_dir + "/cpp_out"; - { - // Create dump dir - std::string cmd = "mkdir -p " + dump_dir; - (void)system(cmd.c_str()); - } - - // Helper to compare intermediate tensors. - // ggml layout [E, W, H], Python ref is [1, H, W, E] (NHWC) - auto compare_nhwc = [&](const std::string & cpp_name, const std::string & ref_name, float tol) { - sam3_dump_state_tensor(*state, cpp_name, dump_dir + "/" + cpp_name); - auto ref_t = load_ref(ref_dir + "/" + ref_name); - auto cpp_t = load_ref(dump_dir + "/" + cpp_name); - if (ref_t.data.empty() || cpp_t.data.empty()) return; - int E = cpp_t.shape[0], nW = cpp_t.shape[1], nH = cpp_t.shape.size() > 2 ? cpp_t.shape[2] : 1; - std::vector transposed(cpp_t.numel()); - for (int e = 0; e < E; ++e) - for (int w = 0; w < nW; ++w) - for (int h = 0; h < nH; ++h) - transposed[h * nW * E + w * E + e] = cpp_t.data[e + w * E + h * E * nW]; - float md = compare_tensors(transposed.data(), ref_t.data.data(), - std::min((int)transposed.size(), ref_t.numel())); - bool ok = md < tol; - fprintf(stderr, " %s %s: max_diff=%.4f\n", ok ? "[PASS]" : "[FAIL]", ref_name.c_str(), md); - if (!ok) n_fail_enc++; - }; - - // Compare intermediate outputs - compare_nhwc("dbg_patch_embed", "patch_embed", 0.05f); - compare_nhwc("dbg_after_ln_pre", "after_ln_pre", 0.1f); - compare_nhwc("dbg_block_0_out", "block_0_out", 0.5f); - - sam3_dump_state_tensor(*state, "vit_output", dump_dir + "/vit_output"); - - // Compare ViT output against Python reference. - // ggml output: [ne0=E, ne1=W, ne2=H] stored as e + w*E + h*E*W - // Python ref: [1, C, H, W] stored as c*H*W + h*W + w (NCHW) - // Need to transpose ggml data to match PyTorch NCHW layout. - auto ref_vit = load_ref(ref_dir + "/vit_output_bchw"); - auto cpp_vit = load_ref(dump_dir + "/vit_output"); - - if (!ref_vit.data.empty() && !cpp_vit.data.empty()) { - // Transpose ggml [E=1024, W=72, H=72] → PyTorch [C=1024, H=72, W=72] - int E = cpp_vit.shape[0], W = cpp_vit.shape[1], H = (cpp_vit.shape.size() > 2) ? cpp_vit.shape[2] : 1; - std::vector transposed(cpp_vit.numel()); - for (int e = 0; e < E; ++e) - for (int h = 0; h < H; ++h) - for (int w = 0; w < W; ++w) - // ggml: e + w*E + h*E*W → pytorch: e*H*W + h*W + w - transposed[e * H * W + h * W + w] = cpp_vit.data[e + w * E + h * E * W]; - - float max_diff = compare_tensors(transposed.data(), ref_vit.data.data(), - std::min((int)transposed.size(), ref_vit.numel())); - bool ok = max_diff < 1.0f; - fprintf(stderr, " %s ViT output vs Python ref: max_diff=%.4f (E=%d W=%d H=%d)\n", - ok ? "[PASS]" : "[FAIL]", max_diff, E, W, H); - if (!ok) n_fail_enc++; - } - - // Dump and compare neck outputs. - // Same layout issue: ggml [C, W, H] vs PyTorch [1, C, H, W] - for (int i = 0; i < 3; ++i) { - std::string name = "neck_det_" + std::to_string(i); - sam3_dump_state_tensor(*state, name, dump_dir + "/" + name); - - auto ref_neck = load_ref(ref_dir + "/" + name); - auto cpp_neck = load_ref(dump_dir + "/" + name); - if (!ref_neck.data.empty() && !cpp_neck.data.empty()) { - int C = cpp_neck.shape[0], nW = cpp_neck.shape[1], nH = (cpp_neck.shape.size() > 2) ? cpp_neck.shape[2] : 1; - std::vector transposed(cpp_neck.numel()); - for (int c = 0; c < C; ++c) - for (int h = 0; h < nH; ++h) - for (int w = 0; w < nW; ++w) - transposed[c * nH * nW + h * nW + w] = cpp_neck.data[c + w * C + h * C * nW]; - - float max_diff = compare_tensors(transposed.data(), ref_neck.data.data(), - std::min((int)transposed.size(), ref_neck.numel())); - bool ok = max_diff < 2.0f; - fprintf(stderr, " %s %s vs Python ref: max_diff=%.4f\n", - ok ? "[PASS]" : "[FAIL]", name.c_str(), max_diff); - } - } - - state.reset(); - sam3_free_model(*model); - model.reset(); - - return true; -} - -// ═══════════════════════════════════════════════════════════════════════════════ -// Main -// ═══════════════════════════════════════════════════════════════════════════════ - -int main(int argc, char ** argv) { - if (argc < 2) { - fprintf(stderr, "Usage: %s [model.ggml] [image.jpg]\n", argv[0]); - fprintf(stderr, "\nRun tests/dump_reference.py first to generate reference data.\n"); - return 1; - } - - std::string ref_dir = argv[1]; - std::string model_path = (argc > 2) ? argv[2] : ""; - std::string image_path = (argc > 3) ? argv[3] : ""; - - int n_pass = 0, n_fail = 0, n_skip = 0; - - // Test RoPE (no model needed) - if (test_rope(ref_dir)) n_pass++; else n_fail++; - - // Test sinusoidal PE (no model needed) - if (test_sinusoidal_pe(ref_dir)) n_pass++; else n_fail++; - - // Test preprocessing (needs image) - if (!image_path.empty()) { - if (test_preprocessing(image_path, ref_dir)) n_pass++; else n_fail++; - } else { - fprintf(stderr, "\n[SKIP] Preprocessing test — no image provided\n"); - n_skip++; - } - - // Test full encoding (needs model + image) - if (!model_path.empty() && !image_path.empty()) { - if (test_encode_image(model_path, image_path, ref_dir)) n_pass++; else n_fail++; - } else { - fprintf(stderr, "\n[SKIP] Full encoding test — no model/image provided\n"); - n_skip++; - } - - fprintf(stderr, "\n═══════════════════════════════════════════\n"); - fprintf(stderr, "Results: %d passed, %d failed, %d skipped\n", n_pass, n_fail, n_skip); - return n_fail > 0 ? 1 : 0; -} diff --git a/tests/trace_single_pixel.py b/tests/trace_single_pixel.py deleted file mode 100644 index e498d6e..0000000 --- a/tests/trace_single_pixel.py +++ /dev/null @@ -1,130 +0,0 @@ -#!/usr/bin/env python3 -"""Trace a single pixel through both C++ and torch bilinear resize.""" -import numpy as np -import torch -from PIL import Image -from torchvision.transforms import v2 - -# Target pixel: (y=36, x=816, c=2) where cpp=110, py=113 - -image_path = "tests/test_random.jpg" -img = Image.open(image_path).convert("RGB") -src = np.array(img, dtype=np.uint8) # [480, 640, 3] HWC -src_h, src_w = src.shape[:2] -img_size = 1008 -c = 2 - -sx = src_w / img_size -sy = src_h / img_size - -y, x = 36, 816 - -# ═══ C++ path (double precision, exact) ═══ -fy = (y + 0.5) * sy - 0.5 -if fy < 0: fy = 0.0 -y0 = int(fy) -y1 = y0 + 1 if y0 < src_h - 1 else y0 -wy = fy - y0 -wy0 = 1.0 - wy - -fx = (x + 0.5) * sx - 0.5 -if fx < 0: fx = 0.0 -x0 = int(fx) -x1 = x0 + 1 if x0 < src_w - 1 else x0 -wx = fx - x0 -wx0 = 1.0 - wx - -p00 = float(src[y0, x0, c]) -p01 = float(src[y0, x1, c]) -p10 = float(src[y1, x0, c]) -p11 = float(src[y1, x1, c]) - -v_double = wy0 * (wx0 * p00 + wx * p01) + wy * (wx0 * p10 + wx * p11) -result_cpp = int(v_double + 0.5) - -print("═══ Pixel Trace: (y=36, x=816, c=2) ═══") -print(f"\nSource image: {src_w}x{src_h}") -print(f"sx={sx:.15f}, sy={sy:.15f}") -print(f"fy={fy:.15f}, fx={fx:.15f}") -print(f"y0={y0}, y1={y1}, wy={wy:.15f}, wy0={wy0:.15f}") -print(f"x0={x0}, x1={x1}, wx={wx:.15f}, wx0={wx0:.15f}") -print(f"p00={p00}, p01={p01}, p10={p10}, p11={p11}") -print(f"v(double)={v_double:.15f}") -print(f"result_cpp = int({v_double:.15f} + 0.5) = {result_cpp}") - -# ═══ torch path ═══ -img_tensor = v2.functional.to_image(img) # uint8 CHW -float_input = img_tensor.float() # [3, 480, 640] - -# Check: what float values does torch see at the source positions? -print(f"\ntorch float input at (c={c}, y0={y0}, x0={x0}): {float_input[c, y0, x0].item()}") -print(f"torch float input at (c={c}, y0={y0}, x1={x1}): {float_input[c, y0, x1].item()}") -print(f"torch float input at (c={c}, y1={y1}, x0={x0}): {float_input[c, y1, x0].item()}") -print(f"torch float input at (c={c}, y1={y1}, x1={x1}): {float_input[c, y1, x1].item()}") - -# Run torch F.interpolate -fi = torch.nn.functional.interpolate( - float_input.unsqueeze(0), size=(img_size, img_size), - mode='bilinear', align_corners=False, antialias=False -).squeeze(0) - -torch_float_result = fi[c, y, x].item() -torch_uint8_result = fi[c, y, x].clamp(0, 255).round().to(torch.uint8).item() - -print(f"\ntorch float result: {torch_float_result:.15f}") -print(f"torch uint8 result: {torch_uint8_result}") -print(f"C++ uint8 result: {result_cpp}") -print(f"difference: {torch_uint8_result - result_cpp}") - -# ═══ Manual torch replication (float32 step by step) ═══ -# Replicate torch's exact computation in float32 -rheight = np.float64(src_h) / np.float64(img_size) -rwidth = np.float64(src_w) / np.float64(img_size) - -h1r = rheight * (y + 0.5) - 0.5 -if h1r < 0: h1r = 0.0 -h1 = int(h1r) -h1p = 1 if h1 < src_h - 1 else 0 -h1lambda = np.float32(h1r - h1) -h0lambda = np.float32(1.0) - h1lambda - -w1r = rwidth * (x + 0.5) - 0.5 -if w1r < 0: w1r = 0.0 -w1 = int(w1r) -w1p = 1 if w1 < src_w - 1 else 0 -w1lambda = np.float32(w1r - w1) -w0lambda = np.float32(1.0) - w1lambda - -# Read as float32 (matching torch's float conversion) -pos1_0 = np.float32(src[h1, w1, c]) -pos1_wp = np.float32(src[h1, w1 + w1p, c]) -pos2_0 = np.float32(src[h1 + h1p, w1, c]) -pos2_wp = np.float32(src[h1 + h1p, w1 + w1p, c]) - -# torch formula: h0lambda * (w0lambda * pos1[0] + w1lambda * pos1[wp]) + h1lambda * (w0lambda * pos2[0] + w1lambda * pos2[wp]) -manual_val = np.float32(h0lambda * (w0lambda * pos1_0 + w1lambda * pos1_wp) + - h1lambda * (w0lambda * pos2_0 + w1lambda * pos2_wp)) - -print(f"\n═══ Manual float32 replication ═══") -print(f"h0lambda={h0lambda:.10f}, h1lambda={h1lambda:.10f}") -print(f"w0lambda={w0lambda:.10f}, w1lambda={w1lambda:.10f}") -print(f"pos: {pos1_0}, {pos1_wp}, {pos2_0}, {pos2_wp}") -print(f"manual float32 result: {manual_val:.10f}") -print(f"manual uint8: {int(round(float(manual_val)))}") - -# ═══ Step by step float32 with explicit intermediate values ═══ -print(f"\n═══ Step-by-step float32 ═══") -t1 = np.float32(w0lambda * pos1_0) -t2 = np.float32(w1lambda * pos1_wp) -row0 = np.float32(t1 + t2) -t3 = np.float32(w0lambda * pos2_0) -t4 = np.float32(w1lambda * pos2_wp) -row1 = np.float32(t3 + t4) -t5 = np.float32(h0lambda * row0) -t6 = np.float32(h1lambda * row1) -result_f32 = np.float32(t5 + t6) -print(f"w0*p00={t1}, w1*p01={t2}, row0={row0}") -print(f"w0*p10={t3}, w1*p11={t4}, row1={row1}") -print(f"h0*row0={t5}, h1*row1={t6}, result={result_f32}") -print(f"uint8(round): {int(round(float(result_f32)))}") -print(f"uint8(+0.5): {int(result_f32 + np.float32(0.5))}") diff --git a/tests/trace_torchvision_resize.py b/tests/trace_torchvision_resize.py deleted file mode 100644 index 1240f65..0000000 --- a/tests/trace_torchvision_resize.py +++ /dev/null @@ -1,124 +0,0 @@ -#!/usr/bin/env python3 -"""Trace exactly what torchvision v2.Resize does on uint8 tensor input. - -The goal: figure out the exact algorithm so C++ can replicate it identically. -""" -import numpy as np -import torch -from torchvision.transforms import v2 -from PIL import Image - -image_path = "tests/test_random.jpg" -img_size = 1008 - -# Load image -img = Image.open(image_path).convert("RGB") -img_tensor = v2.functional.to_image(img) # uint8 CHW - -# ═══ Test 1: Does v2.Resize on uint8 go through PIL or torch? ═══ -# Create a unique pattern that would reveal the algorithm -print("═══ Test 1: PIL vs torch path detection ═══") - -# Check if v2.Resize produces the same output as PIL resize -step1 = v2.ToDtype(torch.uint8, scale=True)(img_tensor) -tv_out = v2.Resize(size=(img_size, img_size))(step1) # uint8 output - -# PIL path: convert to PIL Image, resize, convert back -pil_img = Image.fromarray(step1.permute(1, 2, 0).numpy(), "RGB") -pil_resized = pil_img.resize((img_size, img_size), Image.BILINEAR) -pil_out = torch.from_numpy(np.array(pil_resized)).permute(2, 0, 1) - -diff = (tv_out.to(torch.int16) - pil_out.to(torch.int16)).abs() -print(f" v2.Resize vs PIL.resize(BILINEAR): max_diff={diff.max().item()}, n_diff={(diff>0).sum().item()}") - -# ═══ Test 2: Does v2.Resize use float internally? ═══ -# Resize through float path manually: uint8 → float → interpolate → round → uint8 -print("\n═══ Test 2: Manual float interpolation ═══") -float_input = step1.float() # [3, 480, 640] float32 -float_resized = torch.nn.functional.interpolate( - float_input.unsqueeze(0), # [1, 3, 480, 640] - size=(img_size, img_size), - mode='bilinear', - align_corners=False, - antialias=False, -).squeeze(0) # [3, 1008, 1008] - -# Method A: clamp and truncate -manual_uint8_trunc = float_resized.clamp(0, 255).to(torch.uint8) -diff_a = (tv_out.to(torch.int16) - manual_uint8_trunc.to(torch.int16)).abs() -print(f" vs float→interpolate→clamp→uint8(trunc): max={diff_a.max().item()}, n_diff={(diff_a>0).sum().item()}") - -# Method B: clamp, round, then truncate -manual_uint8_round = float_resized.clamp(0, 255).round().to(torch.uint8) -diff_b = (tv_out.to(torch.int16) - manual_uint8_round.to(torch.int16)).abs() -print(f" vs float→interpolate→clamp→round→uint8: max={diff_b.max().item()}, n_diff={(diff_b>0).sum().item()}") - -# Method C: with antialias=True -float_resized_aa = torch.nn.functional.interpolate( - float_input.unsqueeze(0), - size=(img_size, img_size), - mode='bilinear', - align_corners=False, - antialias=True, -).squeeze(0) -manual_uint8_aa = float_resized_aa.clamp(0, 255).round().to(torch.uint8) -diff_c = (tv_out.to(torch.int16) - manual_uint8_aa.to(torch.int16)).abs() -print(f" vs float→interpolate(aa)→clamp→round→uint8: max={diff_c.max().item()}, n_diff={(diff_c>0).sum().item()}") - -# ═══ Test 3: Check exact matching path ═══ -print("\n═══ Test 3: Finding exact match ═══") - -# Try: PIL → resize directly -# torchvision v2 for uint8 on CPU typically dispatches to PIL -from torchvision.transforms.functional import resize as F_resize -from torchvision.transforms import InterpolationMode - -# F.resize with different options -for aa in [True, False]: - out = F_resize(step1, [img_size, img_size], - interpolation=InterpolationMode.BILINEAR, antialias=aa) - diff = (tv_out.to(torch.int16) - out.to(torch.int16)).abs() - print(f" F.resize(antialias={aa}): max={diff.max().item()}, n_diff={(diff>0).sum().item()}") - -# ═══ Test 4: Save the float interpolation result (before uint8 conversion) ═══ -# This tells us the "ground truth" sub-pixel values -print("\n═══ Test 4: Float interpolation values ═══") -print(f" float_resized range: [{float_resized.min():.4f}, {float_resized.max():.4f}]") -print(f" Sample values at (0,0): {float_resized[:, 0, 0].tolist()}") -print(f" tv_out values at (0,0): {tv_out[:, 0, 0].tolist()}") -print(f" float rounded at (0,0): {float_resized[:, 0, 0].round().tolist()}") - -# ═══ Test 5: What rounding does tv_out use? ═══ -# Compare tv_out against the float values to determine rounding -print("\n═══ Test 5: Rounding analysis ═══") -float_vals = float_resized.flatten().numpy() -tv_vals = tv_out.flatten().numpy().astype(np.float32) - -# For each pixel, check: is tv_val == floor(float_val) or round(float_val)? -n_floor = 0 -n_round = 0 -n_ceil = 0 -n_other = 0 -for i in range(min(len(float_vals), 100000)): - f = float_vals[i] - t = tv_vals[i] - if t == np.floor(f): - n_floor += 1 - elif t == np.round(f): - n_round += 1 - elif t == np.ceil(f): - n_ceil += 1 - else: - n_other += 1 - -print(f" floor: {n_floor}, round: {n_round}, ceil: {n_ceil}, other: {n_other}") -print(f" (sampled first 100000 pixels)") - -# Most likely: v2.Resize on uint8 goes through PIL which uses truncation (floor) -# Let's verify by computing floor explicitly -manual_floor = np.floor(float_vals).astype(np.uint8) -tv_flat = tv_out.flatten().numpy() -floor_match = (manual_floor == tv_flat).sum() -round_match = (np.round(float_vals).astype(np.uint8) == tv_flat).sum() -print(f"\n Global floor match: {floor_match}/{len(tv_flat)} ({100*floor_match/len(tv_flat):.2f}%)") -print(f" Global round match: {round_match}/{len(tv_flat)} ({100*round_match/len(tv_flat):.2f}%)") diff --git a/tests/visualize_cpp_masks.py b/tests/visualize_cpp_masks.py deleted file mode 100644 index 257e207..0000000 --- a/tests/visualize_cpp_masks.py +++ /dev/null @@ -1,65 +0,0 @@ -#!/usr/bin/env python3 -""" -Generate overlay images from C++ SAM2 tracking masks. -Also compares with Python masks side-by-side. -""" -import numpy as np -import os, sys -from PIL import Image - -sam2_root = os.path.expanduser("~/Documents/sam2") -video_dir = os.path.join(sam2_root, "notebooks/videos/bedroom") -py_dir = "/tmp/debug_sam2_pipeline" -cpp_dir = "/tmp/debug_sam2_cpp" -out_dir = os.path.join(os.path.dirname(__file__), "debug_pipeline") -os.makedirs(out_dir, exist_ok=True) - -def load_tensor(path): - data = np.fromfile(path + ".bin", dtype=np.float32) - if os.path.exists(path + ".shape"): - with open(path + ".shape") as f: - shape = [int(x) for x in f.read().strip().split(",") if x] - return data.reshape(shape) if shape else data - return data - -for fi in [1, 2, 5, 10, 20, 30, 40, 50]: - frame_path = os.path.join(video_dir, f"{fi:05d}.jpg") - if not os.path.exists(frame_path): - continue - img = np.array(Image.open(frame_path)) - - # C++ binary mask - cpp_path = os.path.join(cpp_dir, f"f{fi}_output_mask_binary") - if os.path.exists(cpp_path + ".bin"): - cpp_mask = load_tensor(cpp_path) - cpp_binary = cpp_mask > 0.5 - if cpp_binary.shape[0] == img.shape[0] and cpp_binary.shape[1] == img.shape[1]: - overlay_cpp = img.copy() - overlay_cpp[cpp_binary] = (overlay_cpp[cpp_binary] * 0.5 + np.array([255, 0, 0]) * 0.5).astype(np.uint8) - Image.fromarray(overlay_cpp).save(os.path.join(out_dir, f"cpp_frame{fi:03d}.png")) - fg = cpp_binary.sum() - print(f"Frame {fi:3d}: C++ fg={fg:6d} ({100.0*fg/cpp_binary.size:.1f}%)") - - # Python mask (logits) - py_path = os.path.join(py_dir, f"f{fi}_output_mask") - if os.path.exists(py_path + ".bin"): - py_logits = load_tensor(py_path).squeeze() - py_binary = py_logits > 0.0 - if py_binary.shape[0] == img.shape[0] and py_binary.shape[1] == img.shape[1]: - overlay_py = img.copy() - overlay_py[py_binary] = (overlay_py[py_binary] * 0.5 + np.array([0, 0, 255]) * 0.5).astype(np.uint8) - Image.fromarray(overlay_py).save(os.path.join(out_dir, f"py_frame{fi:03d}.png")) - - # IoU comparison - if os.path.exists(cpp_path + ".bin") and os.path.exists(py_path + ".bin"): - cpp_mask2 = load_tensor(cpp_path) - py_logits2 = load_tensor(py_path).squeeze() - c = cpp_mask2.flatten() > 0.5 - p = py_logits2.flatten() > 0.0 - if c.size == p.size: - inter = np.logical_and(c, p).sum() - union = np.logical_or(c, p).sum() - iou = inter / max(union, 1) - print(f" IoU(py,cpp)={iou:.4f}") - -print(f"\nOverlays saved to {out_dir}/")