tracker/scripts/dense_wrapper.py

"""DenseTrackStep: fixed-capacity, dict-free rewrite of SAM3 multiplex track_step
(propagation-only mode) around the REAL submodules, for torch.export / CoreML.

Replaces:
  - output_dict["cond_frame_outputs"/"non_cond_frame_outputs"] dict-walk in
    _prepare_memory_conditioned_features (video_tracking_multiplex.py:1289)
    -> fixed rolling tensor banks (register_buffer) + static concat order
  - MultiplexState add/remove dynamics -> frozen state (1 bucket x multiplex_count);
    its mux/demux matrices become trace-time constants
  - Eager's variable-length skip of missing warmup slots -> fixed banks +
    additive -inf key mask from mem_valid / ptr_valid state buffers

Reuses UNCHANGED: transformer.encoder (memory fusion), _forward_sam_heads
(propagation path incl. MultiplexMaskDecoder), _encode_new_memory
(SimpleMaskEncoder), all with real weights and real dims
(hidden_dim=256, num_maskmem=7, 72x72 memory grid, 1008px masks).

Scope / documented approximations:
  - Exactly 1 cond frame (the init frame). kino's use-case prompts once.
  - Warmup handled via validity masks (not per-count model variants).
  - Cond maskmem tpos follows use_maskmem_tpos_v2 from frame_pos (idx 6 only
    once frame_idx >= num_maskmem).
  - max_objects = multiplex_count = 16 (released SAM3.1 ckpt).
"""

import torch
import torch.nn as nn

import common

E = 72          # sam_image_embedding_size
HW = E * E      # 5184
C = 256
NUM_MASKMEM = 7           # 1 cond + 6 rolling non-cond
N_NONCOND = NUM_MASKMEM - 1
MAX_PTRS = 16             # 1 cond + 15 rolling non-cond
N_PTR = MAX_PTRS - 1
OBJ = common.MULTIPLEX_COUNT  # 16
NUM_OBJ_PTR_TOKENS = MAX_PTRS * OBJ  # 256
# Large finite penalty — CoreML/fp16 friendly vs true -inf
# Softmax key-pad penalty. Keep modest: -1e4 overflowed coremltools' numpy
# attention probe during convert and ballooned RAM to ~60GB.
MASK_NEG = -80.0


class DenseTrackStep(nn.Module):
    def __init__(self, model, ms, pos72):
        super().__init__()
        self.model = model
        self.ms = ms  # frozen MultiplexState: matrices are trace-time constants
        self.multimask = model._use_multimask(False, None)
        assert self.multimask, "propagation multimask expected True"

        # ---- rolling state banks (newest last) ----
        self.register_buffer("mem_bank", torch.zeros(N_NONCOND, 1, C, E, E))
        self.register_buffer("img_bank", torch.zeros(N_NONCOND, HW, 1, C))
        self.register_buffer("ptr_bank", torch.zeros(N_PTR, 1, OBJ, C))
        # 1 = slot filled (attend); 0 = warmup pad (masked)
        self.register_buffer("mem_valid", torch.zeros(N_NONCOND))
        self.register_buffer("ptr_valid", torch.zeros(N_PTR))

        # ---- constants ----
        with torch.no_grad():
            mm_pos = model.maskmem_backbone.position_encoding(
                torch.zeros(1, C, E, E)
            ).to(torch.float32)                      # (1,C,E,E)
        mm_pos = mm_pos.flatten(2).permute(2, 0, 1)  # (HW,1,C)
        tpos = model.maskmem_tpos_enc.detach()       # (7,1,1,C)
        # Non-cond bank: slot s -> t_pos=s+1 -> tpos idx 5-s (fixed).
        # Cond tpos is dynamic from frame_pos (see _cond_tpos).
        noncond_mem_pos = []
        noncond_img_pos = []
        for s in range(N_NONCOND):
            idx = N_NONCOND - 1 - s  # 5..0
            noncond_mem_pos.append(mm_pos + tpos[idx])
            noncond_img_pos.append(pos72 + tpos[idx])
        self.register_buffer("mm_pos", mm_pos.contiguous())
        self.register_buffer("tpos_enc", tpos.contiguous())
        self.register_buffer(
            "noncond_mem_pos", torch.cat(noncond_mem_pos, dim=0).contiguous()
        )
        self.register_buffer(
            "noncond_img_pos", torch.cat(noncond_img_pos, dim=0).contiguous()
        )
        self.register_buffer("pos72", pos72.clone().contiguous())
        self.register_buffer("ptr_tdiffs", torch.arange(1, MAX_PTRS).float())
        self.register_buffer("_one", torch.ones(1))
        self.register_buffer(
            "_arange7", torch.arange(NUM_MASKMEM, dtype=torch.float32)
        )
        self.register_buffer("_oor_idx", torch.tensor(float(NUM_MASKMEM - 1)))

    def _obj_pos(self, frame_pos):
        """Replicates _get_tpos_enc([frame_idx, 1..15], max_abs_pos=16)."""
        from sam3.model.sam3_tracker_utils import get_1d_sine_pe

        pos = torch.cat([frame_pos.reshape(1), self.ptr_tdiffs]) / (MAX_PTRS - 1)
        pe = get_1d_sine_pe(pos, dim=C)                    # (16,C)
        pe = self.model.obj_ptr_tpos_proj(pe)              # (16,C)
        return pe.unsqueeze(1).repeat_interleave(OBJ, dim=0)  # (256,1,C)

    def _cond_tpos(self, frame_pos):
        """use_maskmem_tpos_v2 embedding for cond frame at t=0 (t_pos=frame_idx).

        Soft one-hot over 7 tpos rows — avoids data-dependent integer index
        that torch.export rejects.
        """
        fp = frame_pos.reshape(()).to(self.tpos_enc.dtype)
        in_range = (NUM_MASKMEM - 1) - torch.clamp(fp, min=0, max=float(N_NONCOND))
        sel = torch.where(fp >= float(NUM_MASKMEM), self._oor_idx, in_range)
        # Exact 1 at integer sel, 0 elsewhere (frame_pos is integral in practice)
        w = (1.0 - (self._arange7 - sel).abs()).clamp(min=0.0)
        return (w.view(NUM_MASKMEM, 1, 1, 1) * self.tpos_enc).sum(dim=0)  # (1,1,C)

    def _mem_img_pos_cat(self, frame_pos):
        """Cond spatial+tpos (dynamic) then fixed non-cond bank positions."""
        cond_t = self._cond_tpos(frame_pos)  # (1,1,C)
        cond_mem_pos = self.mm_pos + cond_t
        cond_img_pos = self.pos72 + cond_t
        return (
            torch.cat([cond_mem_pos, self.noncond_mem_pos], dim=0),
            torch.cat([cond_img_pos, self.noncond_img_pos], dim=0),
        )

    def _memory_key_attn_mask(self):
        """Additive SDPA mask (1,1,1,K): 0 attend, MASK_NEG block.

        Key layout matches prompt cat: cond_mem(HW) + mem_bank(N_NONCOND*HW)
        + cond_ptr(OBJ) + flipped ptr_bank(N_PTR*OBJ).
        """
        # mem: cond always valid, then per-slot validity expanded over HW
        mem_slot = torch.cat([self._one, self.mem_valid], dim=0)  # (7,)
        mem_tok = mem_slot.repeat_interleave(HW)  # (7*HW,)
        # ptr: cond always valid; bank is newest-last, prompt uses newest-first
        ptr_slot = torch.cat([self._one, self.ptr_valid.flip(0)], dim=0)  # (16,)
        ptr_tok = ptr_slot.repeat_interleave(OBJ)  # (16*OBJ,)
        valid = torch.cat([mem_tok, ptr_tok], dim=0)  # (K,)
        # (1,1,1,K) broadcasts over batch/heads/queries
        return (MASK_NEG * (1.0 - valid)).view(1, 1, 1, -1)

    def forward(self, vis72, hires0, hires1, cond_mem, cond_img, cond_ptr,
                frame_pos):
        m = self.model

        # ---- dense _prepare_memory_conditioned_features ----
        mem_noncond = (
            self.mem_bank.squeeze(1).flatten(2).permute(0, 2, 1)
            .reshape(N_NONCOND * HW, 1, C)
        )
        mem_tokens = torch.cat(
            [cond_mem.flatten(2).permute(2, 0, 1), mem_noncond], dim=0
        )                                                     # (7*HW,1,C)
        img_tokens = torch.cat(
            [cond_img, self.img_bank.reshape(N_NONCOND * HW, 1, C)], dim=0
        )                                                     # (7*HW,1,C)
        # ptr order: cond, then newest..oldest (t_diff 1..15)
        ptr_noncond = self.ptr_bank.flip(0).squeeze(1).reshape(N_PTR * OBJ, C)
        obj_ptrs = torch.cat(
            [cond_ptr.squeeze(0), ptr_noncond], dim=0
        ).unsqueeze(1)                                        # (256,1,C)

        prompt = torch.cat([mem_tokens, obj_ptrs], dim=0)
        mem_pos_cat, img_pos_cat = self._mem_img_pos_cat(frame_pos)
        prompt_pos = torch.cat([mem_pos_cat, self._obj_pos(frame_pos)], 0)
        key_mask = self._memory_key_attn_mask()

        enc_out = m.transformer.encoder(
            image=vis72,
            src=vis72,
            memory_image=img_tokens,
            memory=prompt,
            image_pos=self.pos72,
            src_pos=self.pos72,
            memory_image_pos=img_pos_cat,
            memory_pos=prompt_pos,
            num_obj_ptr_tokens=NUM_OBJ_PTR_TOKENS,
            memory_key_attn_mask=key_mask,
        )
        pix_feat_with_mem = enc_out["memory"].permute(1, 2, 0).view(1, C, E, E)

        # ---- REAL sam heads (propagation path) ----
        sam_out = m._forward_sam_heads(
            backbone_features=pix_feat_with_mem,
            propagation_high_res_features=[hires0, hires1],
            multimask_output=self.multimask,
            objects_to_interact=list(range(OBJ)),
            multiplex_state=self.ms,
        )
        low_res_masks = sam_out["low_res_masks"]
        high_res_masks = sam_out["high_res_masks"]
        object_score_logits = sam_out["object_score_logits"]
        ious = sam_out["ious"]
        obj_ptr = sam_out["obj_ptr"]

        # ---- REAL memory encoder ----
        maskmem_features, _ = m._encode_new_memory(
            image=None,
            current_vision_feats=[vis72],
            feat_sizes=[(E, E)],
            pred_masks_high_res=high_res_masks,
            object_score_logits=object_score_logits,
            is_mask_from_pts=False,
            conditioning_objects=set(),
            multiplex_state=self.ms,
        )

        # ---- state update: shift + append (static shapes) ----
        self.mem_bank.copy_(
            torch.cat([self.mem_bank[1:], maskmem_features.unsqueeze(0)], 0)
        )
        self.img_bank.copy_(
            torch.cat([self.img_bank[1:], vis72.unsqueeze(0)], 0)
        )
        mux_ptr = self.ms.mux(obj_ptr)  # (1,OBJ,C)
        self.ptr_bank.copy_(
            torch.cat([self.ptr_bank[1:], mux_ptr.unsqueeze(0)], 0)
        )
        self.mem_valid.copy_(torch.cat([self.mem_valid[1:], self._one], 0))
        self.ptr_valid.copy_(torch.cat([self.ptr_valid[1:], self._one], 0))

        return low_res_masks, high_res_masks, object_score_logits, ious


def build_wrapper():
    model = common.build_model()
    pos72 = common.pos_embed_72(model)
    ms = common.make_multiplex_state(model)
    return DenseTrackStep(model, ms, pos72).eval(), model


def seed_banks(wrapper, cache):
    """Fill banks from eager frames so state matches eager just before TARGET."""
    tgt = cache["target_frame"]
    frames = cache["frames"]
    for s in range(N_NONCOND):  # slot s <- frame tgt-6+s
        f = tgt - N_NONCOND + s
        wrapper.mem_bank[s] = frames[f]["maskmem_features"]
        wrapper.img_bank[s] = frames[f]["image_features"]
        wrapper.mem_valid[s] = 1.0
    for p in range(N_PTR):      # slot p <- frame tgt-15+p
        f = tgt - N_PTR + p
        wrapper.ptr_bank[p] = frames[f]["obj_ptr"]
        wrapper.ptr_valid[p] = 1.0


def seed_banks_for_frame(wrapper, cache, frame_idx):
    """Seed banks to match eager state *entering* frame_idx (warmup-aware)."""
    wrapper.mem_bank.zero_()
    wrapper.img_bank.zero_()
    wrapper.ptr_bank.zero_()
    wrapper.mem_valid.zero_()
    wrapper.ptr_valid.zero_()
    frames = cache["frames"]
    # mem: last up to N_NONCOND non-cond frames before frame_idx
    mem_src = [f for f in range(max(1, frame_idx - N_NONCOND), frame_idx)]
    for i, f in enumerate(mem_src):
        slot = N_NONCOND - len(mem_src) + i
        wrapper.mem_bank[slot] = frames[f]["maskmem_features"]
        wrapper.img_bank[slot] = frames[f]["image_features"]
        wrapper.mem_valid[slot] = 1.0
    ptr_src = [f for f in range(max(1, frame_idx - N_PTR), frame_idx)]
    for i, f in enumerate(ptr_src):
        slot = N_PTR - len(ptr_src) + i
        wrapper.ptr_bank[slot] = frames[f]["obj_ptr"]
        wrapper.ptr_valid[slot] = 1.0


def frame_inputs(model, cache, frame_idx):
    feats, _ = common.synth_frame_features(model, frame_idx)
    vis72 = feats[2].contiguous()
    hires0 = feats[0].permute(1, 2, 0).view(1, 32, 4 * E, 4 * E).contiguous()
    hires1 = feats[1].permute(1, 2, 0).view(1, 64, 2 * E, 2 * E).contiguous()
    cond = cache["cond"]
    return (vis72, hires0, hires1,
            cond["maskmem_features"].contiguous(),
            cond["image_features"].contiguous(),
            cond["obj_ptr"].contiguous(),
            torch.tensor([float(frame_idx)]))


def main():
    cache = torch.load("eager_cache.pt", weights_only=False)
    wrapper, model = build_wrapper()

    def report(name, a, b):
        d = (a - b).abs().max().item()
        r = b.abs().max().item()
        print(f"  {name:24s} max|diff|={d:.3e}  (ref max|val|={r:.3e})")
        return d

    # ---- steady-state parity ----
    seed_banks(wrapper, cache)
    tgt = cache["target_frame"]
    inputs = frame_inputs(model, cache, tgt)
    with torch.no_grad():
        low, high, osl, ious = wrapper(*inputs)
    t = cache["target"]
    print(f"parity vs eager track_step @ frame {tgt} (steady):")
    d1 = report("pred_masks(low)", low, t["pred_masks"])
    d2 = report("pred_masks_high_res", high, t["pred_masks_high_res"])
    d3 = report("object_score_logits", osl, t["object_score_logits"])
    d4 = report("ious", ious, t["ious"])
    d5 = report("new maskmem_features", wrapper.mem_bank[-1],
                t["maskmem_features"].squeeze(0))
    d6 = report("new obj_ptr (muxed)", wrapper.ptr_bank[-1], t["obj_ptr"])
    ok_steady = max(d1, d2, d3, d4, d5, d6) < 1e-3

    # ---- warmup parity (frame 3) ----
    warm = cache.get("warmup_frame", 3)
    if "warmup" in cache:
        seed_banks_for_frame(wrapper, cache, warm)
        with torch.no_grad():
            low, high, osl, ious = wrapper(*frame_inputs(model, cache, warm))
        w = cache["warmup"]
        print(f"parity vs eager track_step @ frame {warm} (warmup):")
        w1 = report("pred_masks(low)", low, w["pred_masks"])
        w2 = report("pred_masks_high_res", high, w["pred_masks_high_res"])
        w3 = report("object_score_logits", osl, w["object_score_logits"])
        w4 = report("ious", ious, w["ious"])
        ok_warm = max(w1, w2, w3, w4) < 1e-3
    else:
        print("warmup cache missing — skip warmup check")
        ok_warm = True

    print("PARITY:", "PASS" if (ok_steady and ok_warm) else "FAIL")


if __name__ == "__main__":
    main()