Merge tracker + backbone, drop kino references

README.md

@@ -0,0 +1,101 @@
+---
+license: other
+license_name: sam-license
+license_link: https://github.com/facebookresearch/sam3/blob/main/LICENSE
+library_name: coreml
+pipeline_tag: mask-generation
+base_model: facebook/sam3
+tags:
+  - coreml
+  - sam3
+  - sam3.1
+  - segmentation
+  - video-segmentation
+  - object-tracking
+  - object-multiplex
+  - vision-backbone
+  - stateful
+  - fp16
+  - apple-silicon
+---
+
+# SAM3.1 → CoreML (Apple Silicon)
+
+Real SAM3.1 weights exported to CoreML for on-device video segmentation on Apple Silicon: a stateless per-frame vision backbone plus a stateful multiplex tracker (`track_step`), designed to run together.
+
+- [`backbone/`](backbone) — per-frame image encoder (stateless)
+- [`tracker/`](tracker) — per-frame `track_step` as a stateful CoreML package (multiplex_count=16)
+
+## Contract
+
+The backbone feeds the tracker:
+
+```
+image [1,3,1008,1008] fp32  (normalized to [-1,1])
+  -> vis72  [5184,1,256]
+  -> hires0 [1,32,288,288]   # post sam_mask_decoder.conv_s0
+  -> hires1 [1,64,144,144]   # post conv_s1
+```
+
+Trunk + propagation TriViTDetNeck + `conv_s0`/`conv_s1`. `use_rope_real=True`. Exported via `torch.jit.trace` → coremltools 9 / iOS18. Deploy `CPU_AND_GPU` (~2.4s/frame isolated on Apple Silicon).
+
+## Tracker: `tracker/`
+
+SAM3.1's per-frame `track_step` as a stateful CoreML package on Apple Silicon. Real weights, `multiplex_count=16`, runs on GPU (`CPU_AND_GPU`). Not a product — export + parity proof.
+
+### Layout
+
+- `models/dense_sam3_trackstep.mlpackage` (fp16, 86M) — deploy this. `compute_units=CPU_AND_GPU`.
+- `models/dense_sam3_trackstep_fp32.mlpackage` (fp32, 170M) — precision reference.
+- `scripts/dense_wrapper.py` — rewrites dict memory banks + `MultiplexState` dynamic buckets into fixed `register_buffer`s + validity masks (shift+cat). Warmup = additive `-80` key mask from `mem_valid`/`ptr_valid`. Calls `transformer.encoder`, `_forward_sam_heads`, `_encode_new_memory` as-is once state is frozen.
+- `scripts/common.py` — CPU workarounds, checkpoint download + prefix strip, mux=16.
+- `scripts/eager_baseline.py` — builds `eager_cache.pt` (~388M, not shipped; regenerate once).
+- `scripts/convert_fp16.py`, `export_fp32.py`, `export_coreml.py` — convert. `verify_coreml_lean.py` — low-RAM parity from a saved package.
+- `scripts/toy_bisect.py`, `test_cross_state.py`, `ane_compile_test.py`, `toy_*` — repros for the bugs below.
+- `patches/optimize_state.patch` — required for coremltools 9.0.
+- `triton_stub/` — so `sam3/model/edt.py`'s bare `import triton` doesn't die on Mac (no wheel). The EDT kernel itself is CUDA; track_step export doesn't need it. Swap the call site for `cv2.distanceTransform` / scipy if you want a real CPU path.
+
+### Numbers
+
+Real SAM3.1 tracker weights, 16GB Apple Silicon, 1008px masks, 72×72 memory grid, mux=16.
+
+Dense wrapper vs eager `track_step` (max-abs; values up to ~1e3 → rel ~1e-7–1e-5):
+
+| frame | pred_masks | pred_masks_high_res | object_score_logits | ious | new maskmem | new obj_ptr |
+|---|---|---|---|---|---|---|
+| 16 (steady) | 5.96e-05 | 5.46e-05 | 2.16e-05 | 5.25e-04 | 3.61e-05 | 8.64e-06 |
+| 3 (warmup)  | 9.35e-05 | 8.39e-05 | 1.90e-06 | 7.06e-05 | — | — |
+
+CoreML fp16 vs eager (`CPU_AND_GPU`, frames 1→16 from zero state, `verify_coreml_lean.py`): frame 3 `low_rel=1.7e-5 osl_rel=3.1e-3`, frame 16 `low_rel=7.5e-6 osl_rel=3.1e-3`, mask sign-agree 1.0, zero `appear_flips`, worst rel 3.1e-3. Earlier random-weight fp16 appear-flips don't show up with real logits.
+
+`torch.export` + `coremltools.convert(..., minimum_deployment_target=iOS18)` both fine for fp16/fp32 (fp16 convert ~13s).
+
+### Bugs / limits
+
+**coremltools 9.0** — `optimize_state.py::canonicalize_inplace_pattern` snapshots the op list, then deletes ops while iterating; later states grab a dead `before_op` → `KeyError` in `op_to_node`. Hits anything with ≥3 rolling state buffers via shift+cat (this tracker's mem/img/ptr banks). Torch 2.7.0 pin was a red herring. Patch walks past removed ops. Repro: `toy_bisect.py`. File upstream if you care.
+
+**fp32 predict OOMs on 16GB** — convert is fine; `CPU_ONLY` predict at mux=16 blows past 16GB (masked SDPA over ~36k-token mem × 16 objs × 1008px, fp32 intermediates). fp16 + GPU fits. Need more RAM for fp32 runtime.
+
+**ANE** — every stateful CoreML model fails ANE compile (`ANECCompile` error -14) on macOS 26.5.2 / coremltools 9.0, including a one-buffer + Linear toy. Stateless attention ANE-compiles fine. Platform wall, not this graph. Deploy on `CPU_AND_GPU`: ~0.7s/frame steady, ~7s first call (JIT). Repros in `toy_*` / `ane_compile_test.py`.
+
+### Not wired yet
+
+- One conditioning frame only. Model allows `max_cond_frames_in_attn=4`.
+
+## Backbone: `backbone/`
+
+Per-frame image encoder, stateless. Feeds the tracker above with `vis72` / `hires0` / `hires1` (see contract above).
+
+### Layout
+
+- `models/sam3_vision_backbone.mlpackage` (~875MB fp16)
+
+## Weights & license
+
+`tracker/models/*.mlpackage` and `backbone/models/*.mlpackage` ship **real SAM3.1 weights** (tracker: 457 tensors from `sam3.1_multiplex.pt`, `tracker.model.` prefix stripped; backbone: `detector.backbone.vision_backbone.*`). Meta gates the original on `facebook/sam3`; export scripts pull from the community mirror `AEmotionStudio/sam3.1` (`tracker/scripts/common.py::resolve_checkpoint`, override with `SAM3_HF_REPO` / `SAM3_CHECKPOINT`). That mirror is unofficial — check against Meta's release before you trust it.
+
+Derivative of Meta's SAM Materials → subject to Meta's [SAM License](https://github.com/facebookresearch/sam3/blob/main/LICENSE) (share-alike, field-of-use, attribution). Not a permissive license. Redistributing the weight binaries means honoring that.
+
+## Env
+
+Python 3.9. Torch **2.7.0** + torchvision **0.22.0** (coremltools 9.0 ceiling; `timm` et al. will silently bump torch — check). `coremltools==9.0` with `patches/optimize_state.patch` on `…/optimize_state.py`. Register a `where.scalarother` op at convert (upstream only ships `scalarself`). State inputs `.contiguous()`. `use_rope_real=True` (no complex tensors in CoreML). Drop the triton stub from `sys.path`/`sys.modules` before `torch.export` (`common.hide_triton_stub`).

backbone/models/sam3_vision_backbone.mlpackage/Data/com.apple.CoreML/model.mlmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3c4474e3bf070f097c2a47e5814c33d9fe77bb416bac30011a0e20dd334dfe17
+size 538552

backbone/models/sam3_vision_backbone.mlpackage/Data/com.apple.CoreML/weights/weight.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:94838c0023ea7308bc86782b2f2c124470a3c2d7739961ef2c956c1932a5e178
+size 916622976

backbone/models/sam3_vision_backbone.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "08F5CD35-FEAC-45F5-9D32-511890C3F562": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        },
+        "E4A741A9-FAEE-42AF-893D-1BBFBFD4F9A6": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        }
+    },
+    "rootModelIdentifier": "E4A741A9-FAEE-42AF-893D-1BBFBFD4F9A6"
+}

tracker/models/dense_sam3_trackstep.mlpackage/Data/com.apple.CoreML/model.mlmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b6fcc12ed3815d28f9a271774f71b2e05a02c201f633bbba9cc157ba468930ef
+size 314910

tracker/models/dense_sam3_trackstep.mlpackage/Data/com.apple.CoreML/weights/weight.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:742aba93b515b8ccdc6935c557a8767115538a9fd6009164786f77514ccef6e4
+size 89035968

tracker/models/dense_sam3_trackstep.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "5C239A58-0E5C-41C9-8A56-D0FB690285BA": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        },
+        "DFD8B1C7-D690-4ECD-911C-987F37641012": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        }
+    },
+    "rootModelIdentifier": "DFD8B1C7-D690-4ECD-911C-987F37641012"
+}

tracker/models/dense_sam3_trackstep_fp32.mlpackage/Data/com.apple.CoreML/model.mlmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ee221f2a9b0a2ac1e5798a0462a9e0d038d9e6f3a8d9580363e6f27d7b06f39d
+size 284296

tracker/models/dense_sam3_trackstep_fp32.mlpackage/Data/com.apple.CoreML/weights/weight.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2f57bab04b7ef412f7c3061bfdb1a0a5cdb01165bf61f4fb562e0aa366a1dac5
+size 178052032

tracker/models/dense_sam3_trackstep_fp32.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "86FEDE2A-0538-44EF-9A45-F1DB30FED223": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        },
+        "D1A362A2-9B47-4E58-BB64-CD42D4A1D425": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        }
+    },
+    "rootModelIdentifier": "D1A362A2-9B47-4E58-BB64-CD42D4A1D425"
+}

tracker/patches/optimize_state.patch

@@ -0,0 +1,22 @@
+--- /var/folders/n8/bv7zgpwx3b9crc3mbp6zqb5w0000gn/T/tmp.eWO02SiwSF/coremltools/converters/mil/mil/passes/defs/optimize_state.py	2025-11-03 23:10:10
++++ /private/tmp/claude-501/-Users-aiden-Developer-Kino-kino/54f28767-a77e-430f-881b-85e18da4ab2e/scratchpad/.venv/lib/python3.9/site-packages/coremltools/converters/mil/mil/passes/defs/optimize_state.py	2026-07-23 21:53:13
+@@ -87,7 +87,18 @@
+ 
+         coreml_update_state_ops = self._try_find_child_coreml_update_state_ops(op)
+         for coreml_update_state_op in coreml_update_state_ops:
+-            before_op = block_operation_list[block_operation_list.index(op) + 1]
++            # PATCued: block_operation_list is a stale snapshot; earlier iterations
++            # may have removed ops (their enclosing_block is None). Walk forward to
++            # the next op still present in the block, else insert at block end.
++            idx = block_operation_list.index(op) + 1
++            while (
++                idx < len(block_operation_list)
++                and block_operation_list[idx].enclosing_block is None
++            ):
++                idx += 1
++            before_op = (
++                block_operation_list[idx] if idx < len(block_operation_list) else None
++            )
+             scopes = self._construct_scope_info_list_from_op_scopes(op)
+             with mb.scope(*scopes):
+                 immediate_coreml_update_state = mb.coreml_update_state(

tracker/scripts/ane_compile_test.py

@@ -0,0 +1,61 @@
+"""Step 4b: attempt on-device ANE compile of the converted model.
+Loads dense_sam3_trackstep.mlpackage with CPU_AND_NE and runs one predict.
+The ANECompiler runs at load; failures surface as load errors or CPU fallback.
+"""
+
+import sys
+import time
+import numpy as np
+import torch
+
+import common
+import dense_wrapper as dw
+
+
+def main():
+    import coremltools as ct
+
+    cache = torch.load("eager_cache.pt", weights_only=False)
+    # need model only for input synthesis
+    model = common.build_model()
+    inputs = dw.frame_inputs(model, cache, cache["target_frame"])
+
+    unit = getattr(ct.ComputeUnit, sys.argv[1] if len(sys.argv) > 1 else "CPU_AND_NE")
+    print(f"loading with compute_units={unit} ...")
+    t0 = time.time()
+    mlmodel = ct.models.MLModel("dense_sam3_trackstep.mlpackage",
+                                compute_units=unit)
+    print(f"load OK in {time.time()-t0:.1f}s")
+
+    in_names = [i.name for i in mlmodel.input_description._fd_spec]
+    feed = {n: v.numpy().astype(np.float32) for n, v in zip(in_names, inputs)}
+    state = mlmodel.make_state()
+    t0 = time.time()
+    out = mlmodel.predict(feed, state=state)
+    print(f"predict 1 OK in {time.time()-t0:.1f}s")
+    t0 = time.time()
+    out = mlmodel.predict(feed, state=state)
+    print(f"predict 2 OK in {time.time()-t0:.1f}s")
+    print("outputs:", {k: np.asarray(v).shape for k, v in out.items()})
+
+    # which units actually got used?
+    try:
+        plan = ct.models.compute_plan.MLComputePlan.load_from_path(
+            mlmodel.get_compiled_model_path(), compute_units=unit
+        )
+        counts = {}
+        prog = plan.model_structure.program
+        fn = prog.functions["main"]
+        for op in fn.block.operations:
+            usage = plan.get_compute_device_usage_for_mlprogram_operation(op)
+            if usage is None:
+                continue
+            dev = type(usage.preferred_compute_device).__name__
+            counts[dev] = counts.get(dev, 0) + 1
+        print("op placement:", counts)
+    except Exception as e:
+        print(f"compute plan inspection failed: {type(e).__name__}: {e}")
+
+
+if __name__ == "__main__":
+    main()

tracker/scripts/common.py

@@ -0,0 +1,153 @@
+"""Shared setup: CPU workarounds + model build for SAM3 multiplex video model."""
+
+import sys
+import os
+
+SP = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, os.path.join(SP, "triton_stub"))  # satisfy `import triton`
+sys.path.insert(0, os.path.join(SP, "sam3"))
+
+import torch  # noqa: E402
+
+# CPU-run workarounds: sam3 hardcodes .cuda()/.pin_memory() in eval paths
+torch.Tensor.cuda = lambda self, *a, **kw: self
+torch.Tensor.pin_memory = lambda self, *a, **kw: self
+
+MULTIPLEX_COUNT = 16  # matches released SAM3.1 multiplex checkpoint
+
+
+# Official facebook/sam3.1 is gated; open mirror hosts the same multiplex weights.
+SAM31_HF_REPO = os.environ.get("SAM3_HF_REPO", "AEmotionStudio/sam3.1")
+SAM31_CKPT_NAME = "sam3.1_multiplex.pt"
+
+
+def resolve_checkpoint(checkpoint_path=None):
+    """Return local path to sam3.1_multiplex.pt (download via HF if needed)."""
+    if checkpoint_path:
+        return checkpoint_path
+    env = os.environ.get("SAM3_CHECKPOINT")
+    if env:
+        return env
+    from huggingface_hub import hf_hub_download
+
+    return hf_hub_download(repo_id=SAM31_HF_REPO, filename=SAM31_CKPT_NAME)
+
+
+def hide_triton_stub():
+    """Remove the triton stub after all sam3 modules are imported, so torch
+    inductor's `import triton` fails cleanly into its `triton = None` path."""
+    stub = os.path.join(SP, "triton_stub")
+    while stub in sys.path:
+        sys.path.remove(stub)
+    for k in list(sys.modules):
+        if k == "triton" or k.startswith("triton."):
+            del sys.modules[k]
+
+
+def _tracker_state_dict(ckpt):
+    """Extract tracker weights from merged SAM3.1 multiplex checkpoint.
+
+    HF / mirror file stores full predictor keys (`tracker.model.*` +
+    `detector.*`). Tracker-only module wants the stripped `tracker.model.`
+    prefix. Also accepts a bare tracker state_dict or `{"model": ...}`.
+    """
+    if isinstance(ckpt, dict) and "model" in ckpt and isinstance(ckpt["model"], dict):
+        ckpt = ckpt["model"]
+    if any(k.startswith("tracker.model.") for k in ckpt):
+        return {
+            k[len("tracker.model.") :]: v
+            for k, v in ckpt.items()
+            if k.startswith("tracker.model.")
+        }
+    return ckpt
+
+
+def build_model(
+    multiplex_count=MULTIPLEX_COUNT,
+    checkpoint_path=None,
+    load_checkpoint=True,
+    seed=1234,
+):
+    """Build multiplex video tracker. Default: real SAM3.1 weights, mux=16.
+
+    Set load_checkpoint=False for random-weight debug (uses seed).
+    """
+    from sam3.model_builder import build_sam3_multiplex_video_model
+
+    if not load_checkpoint:
+        torch.manual_seed(seed)
+
+    # Always construct without builder-side load: merged HF ckpt needs key strip.
+    model = build_sam3_multiplex_video_model(
+        checkpoint_path=None,
+        load_from_HF=False,
+        multiplex_count=multiplex_count,
+        use_fa3=False,
+        use_rope_real=True,  # avoid complex tensors (CoreML cannot represent them)
+        device="cpu",
+        strict_state_dict_loading=False,
+    )
+    # Drop image backbone before load: track_step uses precomputed feats; HF
+    # multiplex file also keeps backbone under detector.*, not tracker.model.*.
+    model.backbone = None
+
+    if load_checkpoint:
+        ckpt_path = resolve_checkpoint(checkpoint_path)
+        print(f"loading checkpoint: {ckpt_path}")
+        raw = torch.load(ckpt_path, map_location="cpu", weights_only=True)
+        sd = _tracker_state_dict(raw)
+        missing, unexpected = model.load_state_dict(sd, strict=True)
+        assert not missing and not unexpected, (missing, unexpected)
+        print(f"loaded tracker weights: {len(sd)} tensors")
+
+    model.eval()
+    model.requires_grad_(False)
+    return model
+
+
+def make_multiplex_state(model, num_objects=MULTIPLEX_COUNT):
+    return model.multiplex_controller.get_state(
+        num_objects, torch.device("cpu"), torch.float32, random=False
+    )
+
+
+def synth_frame_features(model, frame_idx, seed=777):
+    """Deterministic synthetic backbone features for one frame (stand-in for the
+    stateless image encoder, which converts separately)."""
+    g = torch.Generator().manual_seed(seed * 100003 + frame_idx)
+    e = model.sam_image_embedding_size  # 72
+    sizes = [(4 * e, 4 * e), (2 * e, 2 * e), (e, e)]
+    # levels 0/1 are pre-projected by sam_mask_decoder.conv_s0/s1 upstream of
+    # track_step (see forward_image), so channels are 32/64, not 256
+    chans = [32, 64, 256]
+    feats = [torch.randn((h * w, 1, c), generator=g) * 0.5
+             for (h, w), c in zip(sizes, chans)]
+    return feats, sizes
+
+
+def pos_embed_72(model):
+    """Constant sine positional embedding of the 72x72 grid, (HW, 1, C)."""
+    from sam3.model.position_encoding import PositionEmbeddingSine
+
+    pe = PositionEmbeddingSine(
+        num_pos_feats=256, normalize=True, scale=None, temperature=10000
+    )
+    e = model.sam_image_embedding_size
+    with torch.no_grad():
+        pos = pe(torch.zeros(1, 1, e, e))  # (1,256,e,e)
+    return pos.flatten(2).permute(2, 0, 1)  # (HW,1,256)
+
+
+def init_masks(num_objects=MULTIPLEX_COUNT, size=1008):
+    """Non-overlapping square blobs for up to 16 objects: (O,1,size,size)."""
+    m = torch.zeros(num_objects, 1, size, size)
+    # 4x4 grid of blobs inside 1008
+    cols = 4
+    cell = size // cols
+    margin = cell // 8
+    for i in range(num_objects):
+        r, c = divmod(i, cols)
+        r0, c0 = r * cell + margin, c * cell + margin
+        r1, c1 = (r + 1) * cell - margin, (c + 1) * cell - margin
+        m[i, 0, r0:r1, c0:c1] = 1.0
+    return m

tracker/scripts/convert_fp16.py

@@ -0,0 +1,52 @@
+"""Convert-only fp16 export of DenseTrackStep (real SAM3.1 weights) and save.
+Splits convert from verify so the RAM-heavy predict loop doesn't OOM the
+convert on a 16GB host. Verify separately with verify_coreml_lean.py.
+"""
+
+import time
+import torch
+
+import common
+import dense_wrapper as dw
+
+
+def main():
+    cache = torch.load("eager_cache.pt", weights_only=False)
+    wrapper, model = dw.build_wrapper()
+    inputs = dw.frame_inputs(model, cache, 3)  # mid frame exercises cond-tpos path
+
+    common.hide_triton_stub()
+    with torch.no_grad():
+        try:
+            ep = torch.export.export(wrapper, inputs)
+        except Exception as e:
+            print(f"strict export failed ({type(e).__name__}); retry strict=False")
+            ep = torch.export.export(wrapper, inputs, strict=False)
+        ep = ep.run_decompositions({})
+    print("torch.export OK")
+
+    import coremltools as ct
+    from coremltools.converters.mil.frontend.torch.torch_op_registry import (
+        register_torch_op,
+    )
+    from coremltools.converters.mil.frontend.torch.ops import _get_inputs
+    from coremltools.converters.mil.mil import Builder as mb
+
+    @register_torch_op(torch_alias=["where.scalarother"])
+    def where_scalarother(context, node):
+        cond, a, b = _get_inputs(context=context, node=node, expected=3)
+        context.add(mb.select(cond=cond, a=a, b=b), node.name)
+
+    t0 = time.time()
+    mlmodel = ct.convert(
+        ep,
+        minimum_deployment_target=ct.target.iOS18,
+        compute_units=ct.ComputeUnit.CPU_ONLY,  # convert/load only; deploy sets CU
+    )
+    print(f"fp16 convert OK in {time.time()-t0:.1f}s")
+    mlmodel.save("dense_sam3_trackstep.mlpackage")
+    print("saved dense_sam3_trackstep.mlpackage")
+
+
+if __name__ == "__main__":
+    main()

tracker/scripts/coreml_traj_test.py

@@ -0,0 +1,52 @@
+"""Trajectory parity: torch wrapper vs converted CoreML model, both from zero
+state, 18 sequential frames with real per-frame inputs. Early frames sit on the
+object-score decision boundary (degenerate zero memory); the test is whether
+late frames track once banks hold real memories.
+"""
+
+import numpy as np
+import torch
+
+import common
+import dense_wrapper as dw
+
+
+def main():
+    import coremltools as ct
+
+    cache = torch.load("eager_cache.pt", weights_only=False)
+    wrapper, model = dw.build_wrapper()
+    wrapper.mem_bank.zero_(); wrapper.img_bank.zero_(); wrapper.ptr_bank.zero_()
+
+    mlmodel = ct.models.MLModel("dense_sam3_trackstep.mlpackage",
+                                compute_units=ct.ComputeUnit.CPU_ONLY)
+    in_names = [i.name for i in mlmodel.input_description._fd_spec]
+    state = mlmodel.make_state()
+
+    frames = list(range(16, 34))
+    print(f"{'frame':>5} {'osl_rel':>9} {'iou_rel':>9} {'mask_sign_agree':>15} "
+          f"{'lowmask_rel':>11}")
+    for f in frames:
+        inputs = dw.frame_inputs(model, cache, f)
+        with torch.no_grad():
+            t_low, t_high, t_osl, t_ious = [x.clone() for x in wrapper(*inputs)]
+        feed = {n: v.numpy().astype(np.float32)
+                for n, v in zip(in_names, inputs)}
+        got = mlmodel.predict(feed, state=state)
+        # map outputs by shape (names are auto-generated)
+        by_shape = {tuple(np.asarray(v).shape): torch.from_numpy(
+            np.asarray(v)).float() for v in got.values()}
+        c_low = by_shape[tuple(t_low.shape)]
+        c_osl = by_shape[tuple(t_osl.shape)]
+        c_ious = by_shape[tuple(t_ious.shape)]
+
+        def rel(a, b):
+            return ((a - b).abs().max() / b.abs().max().clamp_min(1e-9)).item()
+
+        sign_agree = ((c_low > 0) == (t_low > 0)).float().mean().item()
+        print(f"{f:>5} {rel(c_osl, t_osl):>9.2e} {rel(c_ious, t_ious):>9.2e} "
+              f"{sign_agree:>15.4f} {rel(c_low, t_low):>11.2e}")
+
+
+if __name__ == "__main__":
+    main()

tracker/scripts/dense_wrapper.py

@@ -0,0 +1,329 @@
+"""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()

tracker/scripts/eager_baseline.py

@@ -0,0 +1,112 @@
+"""Eager SAM3 track_step baseline: frame 0 mask-init + frames 1..TARGET propagation.
+Caches per-frame outputs needed by the dense wrapper + TARGET/warmup eager
+outputs for parity checking. One-time cost; results go to eager_cache.pt.
+"""
+
+import time
+import torch
+import common
+
+TARGET = 16  # first frame with FULL banks: mem needs f-1..f-6, ptrs need f-1..f-15
+WARMUP = 3   # partial banks — validates validity-mask path
+
+
+def bf_prop(model, frame_idx):
+    feats, sizes = common.synth_frame_features(model, frame_idx)
+    pos72 = common.POS72
+    return {
+        "vision_feats": feats,
+        "vision_masks": [None, None, None],
+        "vision_pos_embeds": [None, None, pos72],
+        "feat_sizes": sizes,
+    }
+
+
+def bf_inter(model, frame_idx):
+    feats, sizes = common.synth_frame_features(model, frame_idx)
+    return {"vision_feats": feats, "feat_sizes": sizes}
+
+
+def main():
+    model = common.build_model()
+    common.POS72 = common.pos_embed_72(model)
+    ms = common.make_multiplex_state(model)
+    output_dict = {"cond_frame_outputs": {}, "non_cond_frame_outputs": {}}
+
+    with torch.no_grad():
+        # ---- frame 0: init via mask input (mode: mask_as_output) ----
+        t0 = time.time()
+        out0 = model.track_step(
+            frame_idx=0,
+            is_init_cond_frame=True,
+            backbone_features_interactive=bf_inter(model, 0),
+            backbone_features_propagation=bf_prop(model, 0),
+            image=None,
+            point_inputs=None,
+            mask_inputs=common.init_masks(),
+            gt_masks=None,
+            frames_to_add_correction_pt=[],
+            output_dict=output_dict,
+            num_frames=64,
+            multiplex_state=ms,
+        )
+        output_dict["cond_frame_outputs"][0] = out0
+        print(f"frame 0 (init): {time.time()-t0:.1f}s")
+        for k, v in out0.items():
+            if torch.is_tensor(v):
+                print(f"  {k}: {tuple(v.shape)} {v.dtype}")
+            elif isinstance(v, list) and v and torch.is_tensor(v[0]):
+                print(f"  {k}: list[{tuple(v[0].shape)}]")
+            else:
+                print(f"  {k}: {type(v).__name__} {v if isinstance(v, set) else ''}")
+
+        # ---- frames 1..TARGET: propagation only ----
+        for f in range(1, TARGET + 1):
+            t0 = time.time()
+            out = model.track_step(
+                frame_idx=f,
+                is_init_cond_frame=False,
+                backbone_features_interactive=None,
+                backbone_features_propagation=bf_prop(model, f),
+                image=None,
+                point_inputs=None,
+                mask_inputs=None,
+                gt_masks=None,
+                frames_to_add_correction_pt=[],
+                output_dict=output_dict,
+                num_frames=64,
+                multiplex_state=ms,
+            )
+            output_dict["non_cond_frame_outputs"][f] = out
+            print(f"frame {f}: {time.time()-t0:.1f}s  "
+                  f"maskmem {tuple(out['maskmem_features'].shape)} "
+                  f"obj_ptr {tuple(out['obj_ptr'].shape)}")
+
+    # ---- cache what the dense wrapper needs ----
+    keep = ["maskmem_features", "maskmem_pos_enc", "obj_ptr",
+            "image_features", "image_pos_enc", "object_score_logits"]
+    cache = {"cond": {}, "frames": {}, "target": {}, "warmup": {}}
+    for k in keep:
+        cache["cond"][k] = out0.get(k)
+    for f in range(1, TARGET):
+        cache["frames"][f] = {k: output_dict["non_cond_frame_outputs"][f].get(k)
+                              for k in keep}
+
+    def pack_target(frame_idx, bucket):
+        tgt = output_dict["non_cond_frame_outputs"][frame_idx]
+        for k in ["pred_masks", "pred_masks_high_res", "object_score_logits",
+                  "maskmem_features", "obj_ptr"]:
+            bucket[k] = tgt[k]
+        bucket["ious"] = tgt["multistep_pred_ious"][-1]
+
+    pack_target(TARGET, cache["target"])
+    cache["target_frame"] = TARGET
+    pack_target(WARMUP, cache["warmup"])
+    cache["warmup_frame"] = WARMUP
+
+    torch.save(cache, "eager_cache.pt")
+    print("saved eager_cache.pt")
+
+
+if __name__ == "__main__":
+    main()

tracker/scripts/export_coreml.py

@@ -0,0 +1,102 @@
+"""torch.export -> coremltools convert (fp16 default) of DenseTrackStep,
+then CoreML-vs-torch parity over warmup→steady stateful frames.
+"""
+
+import time
+import numpy as np
+import torch
+
+import common
+import dense_wrapper as dw
+
+
+def zero_state(wrapper):
+    wrapper.mem_bank.zero_()
+    wrapper.img_bank.zero_()
+    wrapper.ptr_bank.zero_()
+    wrapper.mem_valid.zero_()
+    wrapper.ptr_valid.zero_()
+
+
+def main():
+    cache = torch.load("eager_cache.pt", weights_only=False)
+    wrapper, model = dw.build_wrapper()
+    inputs = dw.frame_inputs(model, cache, 3)
+
+    common.hide_triton_stub()
+    t0 = time.time()
+    with torch.no_grad():
+        try:
+            ep = torch.export.export(wrapper, inputs)
+        except Exception as e:
+            print(f"strict export failed ({type(e).__name__}: {e}); retrying strict=False")
+            ep = torch.export.export(wrapper, inputs, strict=False)
+        ep = ep.run_decompositions({})
+    print(f"torch.export + decompositions OK in {time.time()-t0:.1f}s")
+
+    import coremltools as ct
+    from coremltools.converters.mil.frontend.torch.torch_op_registry import (
+        register_torch_op,
+    )
+    from coremltools.converters.mil.frontend.torch.ops import _get_inputs
+    from coremltools.converters.mil.mil import Builder as mb
+
+    @register_torch_op(torch_alias=["where.scalarother"])
+    def where_scalarother(context, node):
+        cond, a, b = _get_inputs(context=context, node=node, expected=3)
+        context.add(mb.select(cond=cond, a=a, b=b), node.name)
+
+    t0 = time.time()
+    mlmodel = ct.convert(
+        ep,
+        minimum_deployment_target=ct.target.iOS18,
+        compute_units=ct.ComputeUnit.CPU_AND_GPU,
+    )
+    print(f"coremltools convert OK in {time.time()-t0:.1f}s")
+    mlmodel.save("dense_sam3_trackstep.mlpackage")
+    print("saved dense_sam3_trackstep.mlpackage")
+
+    zero_state(wrapper)
+    frames = list(range(1, 19))
+    torch_outs = []
+    with torch.no_grad():
+        for f in frames:
+            torch_outs.append([t.clone() for t in wrapper(*dw.frame_inputs(model, cache, f))])
+
+    in_names = [i.name for i in mlmodel.input_description._fd_spec]
+    out_names = [o.name for o in mlmodel.output_description._fd_spec]
+    print("inputs:", in_names)
+    print("outputs:", out_names)
+    print(f"{'frame':>5} {'osl_rel':>9} {'iou_rel':>9} {'sign_agree':>10} {'low_rel':>9}")
+
+    state = mlmodel.make_state()
+    worst = 0.0
+    appear_flips = 0
+    for i, f in enumerate(frames):
+        feed = {n: v.numpy().astype(np.float32)
+                for n, v in zip(in_names, dw.frame_inputs(model, cache, f))}
+        got = mlmodel.predict(feed, state=state)
+        by_shape = {tuple(np.asarray(v).shape): torch.from_numpy(
+            np.asarray(v)).float() for v in got.values()}
+        t_low, t_high, t_osl, t_ious = torch_outs[i]
+        c_low = by_shape[tuple(t_low.shape)]
+        c_osl = by_shape[tuple(t_osl.shape)]
+        c_ious = by_shape[tuple(t_ious.shape)]
+
+        def rel(a, b):
+            return ((a - b).abs().max() / b.abs().max().clamp_min(1e-9)).item()
+
+        sign = ((c_low > 0) == (t_low > 0)).float().mean().item()
+        r_osl, r_iou, r_low = rel(c_osl, t_osl), rel(c_ious, t_ious), rel(c_low, t_low)
+        worst = max(worst, r_osl, r_iou, r_low)
+        # is_obj_appearing threshold at 0
+        if not torch.equal((c_osl > 0), (t_osl > 0)):
+            appear_flips += 1
+        print(f"{f:>5} {r_osl:>9.2e} {r_iou:>9.2e} {sign:>10.4f} {r_low:>9.2e}")
+
+    print(f"worst rel: {worst:.3e}  appear_flips: {appear_flips}/{len(frames)}")
+    print("FP16 TRAJ:", "PASS" if worst < 5e-2 and appear_flips == 0 else "CHECK")
+
+
+if __name__ == "__main__":
+    main()

tracker/scripts/export_fp32.py

@@ -0,0 +1,91 @@
+"""FLOAT32 CoreML convert + trajectory parity from zero state (warmup→steady)."""
+
+import time
+import numpy as np
+import torch
+
+import common
+import dense_wrapper as dw
+
+
+def zero_state(wrapper):
+    wrapper.mem_bank.zero_()
+    wrapper.img_bank.zero_()
+    wrapper.ptr_bank.zero_()
+    wrapper.mem_valid.zero_()
+    wrapper.ptr_valid.zero_()
+
+
+def main():
+    cache = torch.load("eager_cache.pt", weights_only=False)
+    wrapper, model = dw.build_wrapper()
+    # Export at a mid frame so frame_pos path (dynamic cond tpos) is exercised
+    inputs = dw.frame_inputs(model, cache, 3)
+
+    common.hide_triton_stub()
+    with torch.no_grad():
+        try:
+            ep = torch.export.export(wrapper, inputs)
+        except Exception as e:
+            print(f"strict export failed ({type(e).__name__}: {e}); retrying strict=False")
+            ep = torch.export.export(wrapper, inputs, strict=False)
+        ep = ep.run_decompositions({})
+    print("export OK")
+
+    import coremltools as ct
+    from coremltools.converters.mil.frontend.torch.torch_op_registry import (
+        register_torch_op,
+    )
+    from coremltools.converters.mil.frontend.torch.ops import _get_inputs
+    from coremltools.converters.mil.mil import Builder as mb
+
+    @register_torch_op(torch_alias=["where.scalarother"])
+    def where_scalarother(context, node):
+        cond, a, b = _get_inputs(context=context, node=node, expected=3)
+        context.add(mb.select(cond=cond, a=a, b=b), node.name)
+
+    t0 = time.time()
+    mlmodel = ct.convert(
+        ep,
+        minimum_deployment_target=ct.target.iOS18,
+        # Convert/load on CPU; GPU JIT of masked SDPA previously ballooned RAM.
+        # Deploy path still CPU_AND_GPU via MLModel(..., compute_units=...).
+        compute_units=ct.ComputeUnit.CPU_ONLY,
+        compute_precision=ct.precision.FLOAT32,
+    )
+    print(f"fp32 convert OK in {time.time()-t0:.1f}s")
+    mlmodel.save("dense_sam3_trackstep_fp32.mlpackage")
+
+    zero_state(wrapper)
+    in_names = [i.name for i in mlmodel.input_description._fd_spec]
+    state = mlmodel.make_state()
+    print(f"{'frame':>5} {'osl_rel':>9} {'iou_rel':>9} {'sign_agree':>10} "
+          f"{'lowmask_rel':>11} {'t_nan':>5} {'c_nan':>5}")
+    worst = 0.0
+    for f in range(1, 11):  # warmup→near-steady; enough to prove state+mask
+        fi = dw.frame_inputs(model, cache, f)
+        with torch.no_grad():
+            t_low, t_high, t_osl, t_ious = [x.clone() for x in wrapper(*fi)]
+        feed = {n: v.numpy().astype(np.float32) for n, v in zip(in_names, fi)}
+        got = mlmodel.predict(feed, state=state)
+        by_shape = {tuple(np.asarray(v).shape): torch.from_numpy(
+            np.asarray(v)).float() for v in got.values()}
+        c_low = by_shape[tuple(t_low.shape)]
+        c_osl = by_shape[tuple(t_osl.shape)]
+        c_ious = by_shape[tuple(t_ious.shape)]
+
+        def rel(a, b):
+            return ((a - b).abs().max() / b.abs().max().clamp_min(1e-9)).item()
+
+        sign = ((c_low > 0) == (t_low > 0)).float().mean().item()
+        r_osl, r_iou, r_low = rel(c_osl, t_osl), rel(c_ious, t_ious), rel(c_low, t_low)
+        worst = max(worst, r_osl, r_iou, r_low)
+        print(f"{f:>5} {r_osl:>9.2e} {r_iou:>9.2e} "
+              f"{sign:>10.4f} {r_low:>11.2e} "
+              f"{int(t_low.isnan().sum()):>5} {int(c_low.isnan().sum()):>5}")
+    print(f"worst rel: {worst:.3e}")
+    print("FP32 TRAJ:", "PASS" if worst < 1e-4 else "CHECK")
+
+
+if __name__ == "__main__":
+    main()

tracker/scripts/test_cross_state.py

@@ -0,0 +1,67 @@
+"""Repro of prior spike's crasher: 2 cross-dependent states + attention.
+State A (mem_bank) is READ -> attention -> result WRITTEN to state B (ptr_bank).
+Prior result on torch 2.13 + coremltools 9.0: KeyError in optimize_state.py:93.
+Now testing under torch 2.7.0 (coremltools' tested ceiling).
+"""
+
+import numpy as np
+import torch
+import torch.nn as nn
+import coremltools as ct
+
+C, T, S = 256, 7, 64  # hidden dim, num_maskmem, spatial tokens
+
+
+class CrossStateModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.register_buffer("mem_bank", torch.zeros(T, S, C))  # state A
+        self.register_buffer("ptr_bank", torch.zeros(T, C))     # state B
+        self.attn = nn.MultiheadAttention(C, 4, batch_first=True)
+        self.enc = nn.Linear(C, C)
+
+    def forward(self, x):  # x: (1, S, C)
+        # read BOTH states
+        mem = self.mem_bank.reshape(1, T * S, C)
+        ptrs = self.ptr_bank.reshape(1, T, C)
+        kv = torch.cat([mem, ptrs], dim=1)
+        y, _ = self.attn(x, kv, kv)
+        # write state B from a value derived from reading state A (cross-dependent)
+        new_ptr = y.mean(dim=1)  # (1, C)
+        ptr_updated = torch.cat([self.ptr_bank[1:], new_ptr], dim=0)
+        self.ptr_bank.copy_(ptr_updated)
+        # write state A from current input (shift + cat, static shape)
+        new_mem = self.enc(y)  # (1, S, C)
+        mem_updated = torch.cat([self.mem_bank[1:], new_mem], dim=0)
+        self.mem_bank.copy_(mem_updated)
+        return y
+
+
+def main():
+    torch.manual_seed(0)
+    m = CrossStateModel().eval()
+    m.requires_grad_(False)
+    x = torch.randn(1, S, C)
+
+    with torch.no_grad():
+        ep = torch.export.export(m, (x,))
+        ep = ep.run_decompositions({})
+    print("torch.export OK")
+
+    mlmodel = ct.convert(
+        ep,
+        minimum_deployment_target=ct.target.iOS18,
+        compute_units=ct.ComputeUnit.CPU_ONLY,
+    )
+    print("coremltools convert OK")
+
+    state = mlmodel.make_state()
+    out1 = mlmodel.predict({"x": x.numpy()}, state=state)
+    out2 = mlmodel.predict({"x": x.numpy()}, state=state)
+    k = list(out1)[0]
+    print("predict OK; outputs differ across calls (state evolving):",
+          not np.allclose(out1[k], out2[k]))
+
+
+if __name__ == "__main__":
+    main()

tracker/scripts/toy_ane_test.py

@@ -0,0 +1,17 @@
+"""Does ANY stateful attention model ANE-compile on this machine?
+Reuses the cross-state toy from test_cross_state.py at small scale."""
+import numpy as np, torch, sys
+import coremltools as ct
+from test_cross_state import CrossStateModel, C, S
+
+torch.manual_seed(0)
+m = CrossStateModel().eval(); m.requires_grad_(False)
+x = torch.randn(1, S, C)
+with torch.no_grad():
+    ep = torch.export.export(m, (x,)); ep = ep.run_decompositions({})
+unit = getattr(ct.ComputeUnit, sys.argv[1] if len(sys.argv) > 1 else "CPU_AND_NE")
+ml = ct.convert(ep, minimum_deployment_target=ct.target.iOS18, compute_units=unit)
+st = ml.make_state()
+o1 = ml.predict({"x": x.numpy()}, state=st); o2 = ml.predict({"x": x.numpy()}, state=st)
+k = list(o1)[0]
+print("toy stateful attn on", unit, "OK; state evolves:", not np.allclose(o1[k], o2[k]))

tracker/scripts/toy_bisect.py

@@ -0,0 +1,37 @@
+import numpy as np, torch, torch.nn as nn, sys
+import coremltools as ct
+
+C, S, T = 256, 64, 7
+torch.manual_seed(0)
+
+class Stateless(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.attn = nn.MultiheadAttention(C, 4, batch_first=True)
+        self.enc = nn.Linear(C, C)
+    def forward(self, x):
+        y, _ = self.attn(x, x, x)
+        return self.enc(y)
+
+class StateNoAttn(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.register_buffer("bank", torch.zeros(T, S, C))
+        self.enc = nn.Linear(C, C)
+    def forward(self, x):
+        y = self.enc(x) + self.bank.mean(dim=0, keepdim=True)
+        self.bank.copy_(torch.cat([self.bank[1:], y], 0))
+        return y
+
+which = sys.argv[1]
+m = {"stateless": Stateless, "state_noattn": StateNoAttn}[which]().eval()
+m.requires_grad_(False)
+x = torch.randn(1, S, C)
+with torch.no_grad():
+    ep = torch.export.export(m, (x,)); ep = ep.run_decompositions({})
+ml = ct.convert(ep, minimum_deployment_target=ct.target.iOS18,
+                compute_units=ct.ComputeUnit.CPU_AND_NE)
+st = ml.make_state() if which != "stateless" else None
+kw = {"state": st} if st is not None else {}
+o = ml.predict({"x": x.numpy()}, **kw)
+print(which, "on CPU_AND_NE: predict OK")

tracker/scripts/toy_minstate.py

@@ -0,0 +1,23 @@
+import numpy as np, torch, torch.nn as nn
+import coremltools as ct
+
+torch.manual_seed(0)
+class MinState(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.register_buffer("bank", torch.zeros(1, 64, 256))  # all dims 32-friendly
+        self.enc = nn.Linear(256, 256)
+    def forward(self, x):  # (1,64,256)
+        y = self.enc(x + self.bank)
+        self.bank.copy_(y)  # whole-buffer update, no slice/cat
+        return y
+
+m = MinState().eval(); m.requires_grad_(False)
+x = torch.randn(1, 64, 256)
+with torch.no_grad():
+    ep = torch.export.export(m, (x,)); ep = ep.run_decompositions({})
+ml = ct.convert(ep, minimum_deployment_target=ct.target.iOS18,
+                compute_units=ct.ComputeUnit.CPU_AND_NE)
+st = ml.make_state()
+o = ml.predict({"x": x.numpy()}, state=st)
+print("minimal aligned whole-buffer state on CPU_AND_NE: predict OK")

tracker/scripts/verify_coreml_lean.py

@@ -0,0 +1,70 @@
+"""Low-RAM CoreML runtime parity: load a saved mlpackage, drive it statefully
+frames 1..16 from zero state, compare CoreML outputs to the CACHED eager
+targets (warmup f3, steady f16). No torch-eager forward in the loop and no
+checkpoint load (synth inputs are weight-independent), so peak RAM stays well
+under the 16GB that the in-process double-run in export_fp32.py blows past.
+
+Usage: python verify_coreml_lean.py <mlpackage> [compute_units]
+  compute_units: CPU_ONLY (default) | CPU_AND_GPU
+"""
+
+import sys
+import numpy as np
+import torch
+
+import common
+import dense_wrapper as dw
+
+
+def main():
+    pkg = sys.argv[1] if len(sys.argv) > 1 else "dense_sam3_trackstep_fp32.mlpackage"
+    cu = sys.argv[2] if len(sys.argv) > 2 else "CPU_ONLY"
+
+    cache = torch.load("eager_cache.pt", weights_only=False)
+    # Random-weight model ONLY for weight-independent input generation
+    # (synth features + embedding size). Avoids a second 457-tensor load.
+    model = common.build_model(load_checkpoint=False)
+
+    import coremltools as ct
+    mlmodel = ct.models.MLModel(pkg, compute_units=getattr(ct.ComputeUnit, cu))
+    in_names = [i.name for i in mlmodel.input_description._fd_spec]
+    state = mlmodel.make_state()
+
+    def rel(c, t):
+        t = t.float()
+        c = c.float()
+        return ((c - t).abs().max() / t.abs().max().clamp_min(1e-9)).item()
+
+    checks = {3: cache["warmup"], 16: cache["target"]}
+    worst = 0.0
+    appear_flips = 0
+    print(f"compute_units={cu}  pkg={pkg}")
+    for f in range(1, 17):
+        fi = dw.frame_inputs(model, cache, f)
+        feed = {n: v.numpy().astype(np.float32) for n, v in zip(in_names, fi)}
+        got = mlmodel.predict(feed, state=state)
+        if f in checks:
+            ref = checks[f]
+            by_shape = {tuple(np.asarray(v).shape): torch.from_numpy(
+                np.asarray(v)).float() for v in got.values()}
+            t_low = ref["pred_masks"]
+            t_osl = ref["object_score_logits"]
+            t_iou = ref["ious"]
+            c_low = by_shape[tuple(t_low.shape)]
+            c_osl = by_shape[tuple(t_osl.shape)]
+            c_iou = by_shape[tuple(t_iou.shape)]
+            r_low, r_osl, r_iou = rel(c_low, t_low), rel(c_osl, t_osl), rel(c_iou, t_iou)
+            sign = ((c_low > 0) == (t_low > 0)).float().mean().item()
+            flip = not torch.equal((c_osl > 0), (t_osl > 0))
+            appear_flips += int(flip)
+            worst = max(worst, r_low, r_osl, r_iou)
+            tag = "warmup" if f == 3 else "steady"
+            print(f"frame {f:>2} ({tag}): low_rel={r_low:.2e} osl_rel={r_osl:.2e} "
+                  f"iou_rel={r_iou:.2e} sign_agree={sign:.4f} appear_flip={flip}")
+    thr = 5e-2 if "fp16" in pkg.lower() or cu == "CPU_AND_GPU" else 1e-3
+    print(f"worst rel: {worst:.3e}  appear_flips: {appear_flips}  thr: {thr:.0e}")
+    print("COREML RUNTIME:", "PASS" if worst < thr and appear_flips == 0 else "CHECK")
+
+
+if __name__ == "__main__":
+    main()

tracker/triton_stub/triton/__init__.py

@@ -0,0 +1,12 @@
+def jit(fn=None, **kw):
+    if fn is None:
+        return lambda f: f
+    return fn
+def cdiv(a, b):
+    return (a + b - 1) // b
+class _L:
+    def __getattr__(self, k):
+        return _L()
+    def __call__(self, *a, **k):
+        return _L()
+language = _L()

tracker/triton_stub/triton/language.py

@@ -0,0 +1,12 @@
+class _T:
+    def __getattr__(self, k):
+        return _T()
+    def __call__(self, *a, **k):
+        return _T()
+    def __getitem__(self, k):
+        return _T()
+
+def __getattr__(name):
+    return _T()
+
+constexpr = _T()