slice 12: Faithful SAM3 tracker per-frame track_step (memory-propagation step)

Give the faithful Sam3TrackerVideoModel a real SAM2-style per-frame step built on
the ported slice 8-10 subsystems: memory-attention conditioning (no-memory
first-frame path + 4-layer transformer over stacked maskmem + object-pointer
tokens with temporal PE), prompt encode, mask decode, occlusion-gated object
pointer, and memory encode.

Additive enablers (strictly non-breaking, existing taps unchanged):
- Sam3TrackerMaskDecoder returns the SAM mask token (object-pointer source).
- Sam3TrackerPromptEncoder gains the sparse point/box path + get_dense_pe().

Weight-free component test covers shapes/finiteness/determinism on a tiny grid for
both the init and memory-conditioned paths. Numeric parity stays the deferred
out-of-sandbox gate (parity-status sam3_video). First slice of the phased video
streaming port (slices 12-16).

Author: benjipeng

src/mlx_cv/models/sam3/real_tracker_decoder.py

@@ -99,6 +99,32 @@ def __call__(self, input_masks: mx.array | None, batch_size: int = 1) -> mx.arra
         dense = self.no_mask_embed.weight.reshape(1, 1, 1, -1)
         return mx.broadcast_to(dense, (batch_size, height, width, self.hidden_size))
 
+    def get_dense_pe(self) -> mx.array:
+        """Positional encoding over the image-embedding grid -> ``[1, H, W, hidden]`` (NHWC)."""
+        height, width = self.image_embedding_size
+        ys = (mx.arange(height).astype(mx.float32) + 0.5) / height
+        xs = (mx.arange(width).astype(mx.float32) + 0.5) / width
+        grid = mx.stack(
+            [mx.broadcast_to(xs[None, :], (height, width)), mx.broadcast_to(ys[:, None], (height, width))],
+            axis=-1,
+        )  # [H, W, 2], already normalized -> shared_embedding treats it as such (input_shape=None)
+        return self.shared_embedding(grid)[None]
+
+    def encode_sparse(self, points: mx.array, labels: mx.array) -> mx.array:
+        """Embed point / box-corner prompts into sparse embeddings ``[B, N, hidden]``.
+
+        ``labels``: -1 padding (not-a-point), 0/1 negative/positive point, 2/3 box top-left/
+        bottom-right corner. Mirrors the SAM2 prompt encoder's point path.
+        """
+        coords = points + 0.5  # shift to pixel centers
+        embeddings = self.shared_embedding(coords, input_shape=(self.input_image_size, self.input_image_size))
+        labels_e = labels[..., None]
+        point_embed = self.point_embed.weight  # [num_point_embeddings, hidden]
+        embeddings = mx.where(labels_e == -1, mx.zeros_like(embeddings) + self.not_a_point_embed.weight, embeddings)
+        for label_value in range(point_embed.shape[0]):
+            embeddings = mx.where(labels_e == label_value, embeddings + point_embed[label_value], embeddings)
+        return embeddings
+
 
 class Sam3TrackerAttention(nn.Module):
     """Attention with optional internal downsampling (q/k/v -> internal, o -> hidden)."""
@@ -194,6 +220,7 @@ class Sam3TrackerMaskDecoderOutput:
     masks: mx.array
     iou_pred: mx.array
     object_score_logits: mx.array
+    sam_tokens_out: mx.array  # [B, point_batch, num_masks, C] — mask-token outputs (obj-pointer source)
 
 
 class Sam3TrackerMaskDecoder(nn.Module):
@@ -271,4 +298,10 @@ def __call__(
         mask_slice = slice(1, None) if multimask_output else slice(0, 1)
         masks = masks[:, :, mask_slice]
         iou_pred = iou_pred[:, :, mask_slice]
-        return Sam3TrackerMaskDecoderOutput(masks=masks, iou_pred=iou_pred, object_score_logits=object_score_logits)
+        sam_tokens_out = mask_tokens_out[:, :, mask_slice]
+        return Sam3TrackerMaskDecoderOutput(
+            masks=masks,
+            iou_pred=iou_pred,
+            object_score_logits=object_score_logits,
+            sam_tokens_out=sam_tokens_out,
+        )

src/mlx_cv/models/sam3/real_video_model.py

@@ -15,6 +15,10 @@
 
 from __future__ import annotations
 
+import math
+from dataclasses import dataclass, field
+from typing import Any
+
 import mlx.core as mx
 import mlx.nn as nn
 
@@ -25,7 +29,41 @@
 from .real_video_config import Sam3TrackerVideoConfig
 from .real_vision import Sam3VisionConfig, Sam3VisionNeck
 
-__all__ = ["Sam3TrackerVideoModel", "Sam3VideoModel", "build_sam3_video_real"]
+__all__ = [
+    "Sam3TrackerStageOutput",
+    "Sam3TrackerVideoModel",
+    "Sam3VideoModel",
+    "build_sam3_video_real",
+]
+
+
+def _upsample_nhwc(x: mx.array, factor: int) -> mx.array:
+    """Nearest-neighbour integer upsample of an NHWC tensor along H and W."""
+    if factor == 1:
+        return x
+    return mx.repeat(mx.repeat(x, factor, axis=1), factor, axis=2)
+
+
+def _get_1d_sine_pe(positions: mx.array, dim: int, temperature: float = 10000.0) -> mx.array:
+    """1D sinusoidal positional encoding of ``positions`` -> ``[..., dim]`` (SAM2 ``get_1d_sine_pe``)."""
+    half = dim // 2
+    dim_t = temperature ** (2.0 * (mx.arange(half).astype(mx.float32) // 1) / dim)
+    angles = positions[..., None] / dim_t
+    return mx.concatenate([mx.sin(angles), mx.cos(angles)], axis=-1)
+
+
+@dataclass
+class Sam3TrackerStageOutput:
+    """One frame's faithful tracker output (object batch ``B`` along axis 0)."""
+
+    low_res_masks: mx.array  # NHWC [B, 4g, 4g, 1]
+    high_res_masks: mx.array  # NHWC [B, 16g, 16g, 1]
+    iou_pred: mx.array  # [B, 1, 1]
+    object_score_logits: mx.array  # [B, 1, 1]
+    obj_ptr: mx.array  # [B, hidden_dim]
+    maskmem_features: mx.array | None = None  # NHWC [B, g, g, mem_dim]
+    maskmem_pos_enc: mx.array | None = None  # NHWC [B, g, g, mem_dim]
+    extra: dict[str, Any] = field(default_factory=dict)
 
 
 class Sam3TrackerVideoModel(nn.Module):
@@ -50,6 +88,138 @@ def __init__(self, config: Sam3TrackerVideoConfig, hidden_dim: int = 256, mem_di
         self.no_object_pointer = mx.zeros((1, hidden_dim))
         self.occlusion_spatial_embedding_parameter = mx.zeros((1, mem_dim))
 
+        self.hidden_dim = hidden_dim
+        self.mem_dim = mem_dim
+        self.num_maskmem = config.num_maskmem
+
+    # ---- faithful per-frame streaming step (slice 12) ------------------------
+
+    def _assemble_memory(
+        self, previous_frames: list[Sam3TrackerStageOutput], batch: int
+    ) -> tuple[mx.array, mx.array, int]:
+        """Stack prior-frame spatial memory + object pointers -> ``([mem_seq,B,mem_dim], pos, num_ptr_tokens)``.
+
+        Spatial memory carries the learned ``memory_temporal_positional_encoding`` per slot; object
+        pointers are split into ``hidden_dim // mem_dim`` tokens (appended last so memory-attention can
+        exclude them from RoPE) with a sine temporal encoding projected to ``mem_dim``.
+        """
+        recent = previous_frames[-(self.num_maskmem - 1) :]
+        spatial_mem: list[mx.array] = []
+        spatial_pos: list[mx.array] = []
+        for offset, frame in enumerate(reversed(recent)):
+            t_pos = offset + 1  # 1 == most recent prior frame
+            feats = frame.maskmem_features
+            b, h, w, c = feats.shape
+            spatial_mem.append(feats.reshape(b, h * w, c).transpose(1, 0, 2))
+            pos = frame.maskmem_pos_enc.reshape(b, h * w, c).transpose(1, 0, 2)
+            spatial_pos.append(pos + self.memory_temporal_positional_encoding[self.num_maskmem - t_pos - 1])
+
+        obj_tokens: list[mx.array] = []
+        obj_token_pos: list[mx.array] = []
+        num_obj_ptr_tokens = 0
+        if recent:
+            stacked = mx.stack([frame.obj_ptr for frame in recent], axis=0)  # [n, B, hidden_dim]
+            n = stacked.shape[0]
+            tokens_per_ptr = self.hidden_dim // self.mem_dim
+            stacked = stacked.reshape(n, batch, tokens_per_ptr, self.mem_dim)
+            stacked = stacked.transpose(0, 2, 1, 3).reshape(n * tokens_per_ptr, batch, self.mem_dim)
+            rel = mx.arange(n).astype(mx.float32)
+            sine = self.temporal_positional_encoding_projection_layer(_get_1d_sine_pe(rel, self.hidden_dim))
+            sine = mx.repeat(sine, tokens_per_ptr, axis=0)
+            obj_tokens.append(stacked)
+            obj_token_pos.append(mx.broadcast_to(sine[:, None, :], (n * tokens_per_ptr, batch, self.mem_dim)))
+            num_obj_ptr_tokens = n * tokens_per_ptr
+
+        memory = mx.concatenate(spatial_mem + obj_tokens, axis=0)
+        memory_pos = mx.concatenate(spatial_pos + obj_token_pos, axis=0)
+        return memory, memory_pos, num_obj_ptr_tokens
+
+    def _condition_on_memory(
+        self,
+        vision_features: mx.array,
+        vision_pos: mx.array,
+        previous_frames: list[Sam3TrackerStageOutput] | None,
+        is_init_cond_frame: bool,
+    ) -> mx.array:
+        """Fuse current NHWC features with the memory bank -> NHWC ``[B, g, g, C]``."""
+        batch, height, width, channels = vision_features.shape
+        if is_init_cond_frame or not previous_frames:
+            # First frame: directly add the no-memory embedding (no transformer pass).
+            return vision_features + self.no_memory_embedding.reshape(1, 1, 1, channels)
+
+        memory, memory_pos, num_obj_ptr_tokens = self._assemble_memory(previous_frames, batch)
+        current = vision_features.reshape(batch, height * width, channels).transpose(1, 0, 2)
+        current_pos = vision_pos.reshape(batch, height * width, channels).transpose(1, 0, 2)
+        conditioned = self.memory_attention(
+            current_vision_features=current,
+            memory=memory,
+            current_vision_position_embeddings=current_pos,
+            memory_position_embeddings=memory_pos,
+            num_object_pointer_tokens=num_obj_ptr_tokens,
+        )  # [1, B, seq, C]
+        return conditioned.reshape(batch, height, width, channels)
+
+    def track_step(
+        self,
+        *,
+        vision_features: mx.array,  # NHWC [B, g, g, C]
+        vision_pos: mx.array,  # NHWC [B, g, g, C]
+        high_res_features: list[mx.array] | tuple[mx.array, mx.array],  # [4g-res, 2g-res] raw FPN, NHWC
+        is_init_cond_frame: bool = False,
+        point_inputs: tuple[mx.array, mx.array] | None = None,  # (coords [B,N,2], labels [B,N])
+        mask_inputs: mx.array | None = None,  # NHWC [B, Hm, Wm, 1]
+        previous_frames: list[Sam3TrackerStageOutput] | None = None,
+        run_mem_encoder: bool = True,
+    ) -> Sam3TrackerStageOutput:
+        """One faithful SAM2-style tracker step over a single frame (object batch ``B``)."""
+        batch, grid_h, grid_w, channels = vision_features.shape
+
+        conditioned = self._condition_on_memory(vision_features, vision_pos, previous_frames, is_init_cond_frame)
+
+        # High-res guides for mask upscaling (conv_s0 / conv_s1 own the channel projections).
+        feat_s0 = self.mask_decoder.conv_s0(high_res_features[0])
+        feat_s1 = self.mask_decoder.conv_s1(high_res_features[1])
+
+        # Sparse prompt: explicit points/box, else a single padding point (label -1).
+        if point_inputs is not None:
+            coords, labels = point_inputs
+        else:
+            coords = mx.zeros((batch, 1, 2))
+            labels = -mx.ones((batch, 1))
+        sparse = self.prompt_encoder.encode_sparse(coords, labels)[:, None]  # [B, 1, N, C]
+
+        # Dense prompt + dense positional encoding.
+        dense = self.prompt_encoder(mask_inputs, batch_size=batch)  # NHWC [B, g, g, C]
+        image_pe = mx.broadcast_to(self.prompt_encoder.get_dense_pe(), (batch, grid_h, grid_w, channels))
+
+        decoded = self.mask_decoder(conditioned, image_pe, sparse, dense, [feat_s0, feat_s1], multimask_output=False)
+        low_res_masks = decoded.masks[:, 0, 0][..., None]  # NHWC [B, 4g, 4g, 1]
+        high_res_masks = _upsample_nhwc(low_res_masks, 4)  # NHWC [B, 16g, 16g, 1] for the memory encoder
+
+        # Object pointer from the SAM token, gated by object presence (occlusion handling).
+        sam_token = decoded.sam_tokens_out[:, 0, 0]  # [B, C]
+        obj_ptr = self.object_pointer_proj(sam_token)
+        is_obj = (decoded.object_score_logits[:, 0] > 0).astype(obj_ptr.dtype)  # [B, 1]
+        obj_ptr = is_obj * obj_ptr + (1.0 - is_obj) * self.no_object_pointer
+
+        maskmem_features = maskmem_pos_enc = None
+        if run_mem_encoder:
+            mask_for_mem = mx.sigmoid(high_res_masks)
+            mask_for_mem = mask_for_mem * self.config.sigmoid_scale_for_mem_enc + self.config.sigmoid_bias_for_mem_enc
+            mem = self.memory_encoder(conditioned, mask_for_mem)
+            maskmem_features = mem.vision_features
+            maskmem_pos_enc = mem.vision_pos_enc
+
+        return Sam3TrackerStageOutput(
+            low_res_masks=low_res_masks,
+            high_res_masks=high_res_masks,
+            iou_pred=decoded.iou_pred,
+            object_score_logits=decoded.object_score_logits,
+            obj_ptr=obj_ptr,
+            maskmem_features=maskmem_features,
+            maskmem_pos_enc=maskmem_pos_enc,
+        )
+
 
 class Sam3VideoModel(nn.Module):
     """Faithful SAM 3 video model: detector + tracker + tracker neck (1797 tensors)."""

tests/test_sam3_video_streaming_step_real.py

@@ -0,0 +1,104 @@
+"""Slice 12: faithful SAM3 tracker per-frame ``track_step`` (single object).
+
+Weight-free structural verification of the per-frame memory-propagation step that
+``Sam3TrackerVideoModel`` now exposes (mirrors the SAM2-style tracker loop with the
+faithful slice 8-10 subsystems):
+
+- the init/no-memory path (seed box prompt) produces coherent mask / pointer / memory
+  shapes and finite values;
+- the memory-conditioned path (a prior frame's memory + object pointer fed back through
+  the 4-layer memory-attention transformer) runs and stays finite;
+- the step is deterministic (identical inputs -> identical outputs).
+
+Numeric parity vs the upstream reference is the deferred, out-of-sandbox gate
+(see parity-status ``sam3_video``); no synthetic pass and no real weights here.
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+import pytest
+
+from mlx_cv.models.sam3.real_tracker_decoder import Sam3TrackerPromptEncoderConfig
+from mlx_cv.models.sam3.real_video_config import Sam3TrackerVideoConfig
+from mlx_cv.models.sam3.real_video_model import Sam3TrackerStageOutput, Sam3TrackerVideoModel
+
+GRID = 4  # tracker feature grid (g); image-embedding + RoPE feat sizes matched to it
+CHANNELS = 256
+MEM_DIM = 64
+
+
+def _tiny_tracker() -> Sam3TrackerVideoModel:
+    """A faithful tracker on a tiny grid: prompt-encoder grid and RoPE feat sizes == GRID."""
+    config = Sam3TrackerVideoConfig(
+        prompt_encoder=Sam3TrackerPromptEncoderConfig(image_size=GRID * 16, patch_size=16),
+        memory_attention_rope_feat_sizes=(GRID, GRID),
+    )
+    model = Sam3TrackerVideoModel(config)
+    mx.eval(model.parameters())
+    return model
+
+
+def _frame_features(seed: int) -> dict:
+    """Random per-frame inputs: top-level features + pos enc + two raw high-res FPN levels."""
+    key = mx.random.key(seed)
+    k0, k1, k2, k3 = mx.random.split(key, 4)
+    return {
+        "vision_features": mx.random.normal((1, GRID, GRID, CHANNELS), key=k0),
+        "vision_pos": mx.random.normal((1, GRID, GRID, CHANNELS), key=k1),
+        "high_res_features": [
+            mx.random.normal((1, GRID * 4, GRID * 4, CHANNELS), key=k2),  # 4g-res
+            mx.random.normal((1, GRID * 2, GRID * 2, CHANNELS), key=k3),  # 2g-res
+        ],
+    }
+
+
+def _box_prompt() -> tuple[mx.array, mx.array]:
+    # SAM box encoding: top-left (label 2) + bottom-right (label 3), in input-image pixels.
+    coords = mx.array([[[8.0, 8.0], [40.0, 40.0]]])  # [1, 2, 2]
+    labels = mx.array([[2.0, 3.0]])  # [1, 2]
+    return coords, labels
+
+
+def _assert_stage_shapes(out: Sam3TrackerStageOutput) -> None:
+    assert out.low_res_masks.shape == (1, GRID * 4, GRID * 4, 1)
+    assert out.high_res_masks.shape == (1, GRID * 16, GRID * 16, 1)
+    assert out.iou_pred.shape == (1, 1, 1)
+    assert out.object_score_logits.shape == (1, 1, 1)
+    assert out.obj_ptr.shape == (1, CHANNELS)
+    assert out.maskmem_features.shape == (1, GRID, GRID, MEM_DIM)
+    assert out.maskmem_pos_enc.shape == (1, GRID, GRID, MEM_DIM)
+    for value in (out.low_res_masks, out.high_res_masks, out.obj_ptr, out.maskmem_features):
+        assert bool(mx.all(mx.isfinite(value)).item())
+
+
+def test_init_frame_box_prompt_produces_coherent_shapes():
+    model = _tiny_tracker()
+    out = model.track_step(is_init_cond_frame=True, point_inputs=_box_prompt(), **_frame_features(0))
+    _assert_stage_shapes(out)
+
+
+def test_memory_conditioned_frame_runs_and_is_finite():
+    model = _tiny_tracker()
+    init = model.track_step(is_init_cond_frame=True, point_inputs=_box_prompt(), **_frame_features(0))
+    tracked = model.track_step(is_init_cond_frame=False, previous_frames=[init], **_frame_features(1))
+    _assert_stage_shapes(tracked)
+
+
+def test_track_step_is_deterministic():
+    model = _tiny_tracker()
+    features = _frame_features(0)
+    first = model.track_step(is_init_cond_frame=True, point_inputs=_box_prompt(), **features)
+    second = model.track_step(is_init_cond_frame=True, point_inputs=_box_prompt(), **features)
+    assert bool(mx.all(first.low_res_masks == second.low_res_masks).item())
+    assert bool(mx.all(first.obj_ptr == second.obj_ptr).item())
+    assert bool(mx.all(first.maskmem_features == second.maskmem_features).item())
+
+
+def test_decoder_exposes_sam_token_for_object_pointer():
+    """The additive mask-token return must align with the (single) selected mask."""
+    model = _tiny_tracker()
+    out = model.track_step(is_init_cond_frame=True, point_inputs=_box_prompt(), **_frame_features(2))
+    # obj_ptr is finite and shaped from the SAM token via object_pointer_proj.
+    assert out.obj_ptr.shape == (1, CHANNELS)
+    assert bool(mx.all(mx.isfinite(out.obj_ptr)).item())