MPS-EfficientTAM

A Swift 5.9 package implementing EfficientTAM Tiny 512 image prompting directly with Metal Performance Shaders Graph.

The package includes separate, composable GPU stages for image encoding, point/box decoding, iterative mask prompting, official mask resizing, and the video memory path (memory encoder, memory attention, mask selection). A single-object, forward-only EfficientTAMVideoTracker composes them into a GPU-resident video tracker.

Current scope

• Official efficienttam_ti_512x512 image encoder architecture and weights.
• 12 ViT blocks with the upstream window/global-attention schedule.
• Upstream ImageNet preprocessing and learned no_mem_embed addition.
• Point and box prompt embeddings, two-way transformer, mask hypernetworks, three 128x128 mask-logit outputs, and predicted IoU scores.
• Iterative prompting from one prior 128x128 mask-logit plane through the official learned mask-downscaling path.
• Official bilinear resizing of raw mask logits to a caller-selected output size (align_corners=false), kept separate from decoding.
• Video tracking: object-presence score, object pointers with the no_obj_ptr substitution, NO_OBJ_SCORE mask suppression, the four-stage memory encoder, and the four-layer RoPE memory attention over up to seven spatial memories and sixteen object pointers.
• Float32 transformer intermediates for accuracy. EfficientTAM's image encoder has known severe accuracy degradation when indiscriminately converted to FP16.
• A synchronous run, asynchronous submit, and GPU-resident encode API matching the other Fabric MPS packages.
• No CPU/GPU synchronization in encode; in-flight saturation drops work by returning false rather than blocking.
• Full-tensor accuracy testing against an output generated by the official PyTorch implementation's Apple-MPS behavior.

The image encoder outputs a float32 NCHW embedding with shape [1, 256, 32, 32]. Cache it and reuse it with EfficientTAMPromptDecoder. Every GPU stage accepts caller-owned buffers and command buffers, so clients can schedule each stage independently or compose them into a larger graph.

Requirements

• macOS 15, iOS 18, or visionOS 2
• Swift 5.9
• A Metal device

Usage

import Metal
import MetalPerformanceShaders
import MPSEfficientTAM

let device = MTLCreateSystemDefaultDevice()!
let commandQueue = device.makeCommandQueue()!
let encoder = try EfficientTAMImageEncoder(commandQueue: commandQueue)

let commandBuffer = MPSCommandBuffer(from: commandQueue)
let output = device.makeBuffer(
    length: encoder.outputBufferLength,
    options: .storageModePrivate
)!

let accepted = try encoder.encode(
    inputBuffer: rgbFloat32NHWCBuffer,
    outputBuffer: output,
    commandBuffer: commandBuffer
)
commandBuffer.commit()

let decoder = try EfficientTAMPromptDecoder(promptCount: 2, commandQueue: commandQueue)
let prediction = try decoder.run(
    imageEmbeddingBuffer: output,
    prompts: [
        .init(x: 160, y: 240, label: .positivePoint),
        .init(x: 0, y: 0, label: .padding),
    ]
)

For iterative refinement, encode a selected prior mask and pass its dense embedding to the decoder:

let maskPromptEncoder = try EfficientTAMMaskPromptEncoder(commandQueue: commandQueue)
let densePrompt = device.makeBuffer(
    length: maskPromptEncoder.outputBufferLength,
    options: .storageModePrivate
)!
let maskPromptCommandBuffer = MPSCommandBuffer(from: commandQueue)

_ = try maskPromptEncoder.encode(
    maskLogitsBuffer: prior128x128Logits,
    densePromptEmbeddingBuffer: densePrompt,
    commandBuffer: maskPromptCommandBuffer
)
maskPromptCommandBuffer.commit()

let refinedPrediction = try decoder.run(
    imageEmbeddingBuffer: output,
    prompts: prompts,
    densePromptEmbeddingBuffer: densePrompt
)

EfficientTAMMaskPostprocessor resizes one or more 128x128 mask-logit planes without thresholding them. binaryMask(from:threshold:) is available after CPU readback, while GPU clients can keep the logits resident for their own filtering or compositing.

The package deliberately does not include guided-image or joint-bilateral refinement. Those are useful downstream composition choices, but are not part of EfficientTAM's core learned mask path. Fabric can expose them as separate nodes without coupling third-party package users to Fabric.

inputBuffer is tightly packed float32 RGB in NHWC order, 512x512, with values in 0...1. The model performs ImageNet normalization internally.

run(inputBuffer:) is an offline/testing convenience that returns a Swift array and therefore waits for GPU completion. Real-time consumers should use encode or submit.

Video tracking

EfficientTAMVideoTracker tracks one object forward through a sequence from a prompted first frame. Every submission chains image encode, memory attention, tracking decode, best-IoU mask selection and memory encode on the caller's queue with no CPU/GPU wait, and returns GPU-resident buffers.

let tracker = try EfficientTAMVideoTracker(commandQueue: commandQueue, maxFramesInFlight: 3)
try tracker.prewarmMemoryAttention()   // optional, avoids a first-use compile hitch

// Inside a frame that already has an MPSCommandBuffer (for example Fabric's):
// your own GPU work that produces `modelInput` is encoded first, then the tracker,
// then whatever consumes its output. The frame's owner commits once.
let first = try tracker.encodeInitialFrame(
    inputBuffer: modelInput,
    prompts: [
        .init(x: 160, y: 300, label: .positivePoint),
        .init(x: 0, y: 0, label: .padding),
    ],
    commandBuffer: frameCommandBuffer
)

// Later frames: returns nil (a dropped frame) when maxFramesInFlight are busy.
if let output = try tracker.encodeNextFrame(inputBuffer: modelInput, commandBuffer: frameCommandBuffer)
{
    // output.maskLogitsBuffer        [128, 128] raw logits, feed EfficientTAMMaskPostprocessor
    // output.objectScoreLogitBuffer  object-presence logit; <= 0 means occluded/absent
    // output.iouPredictionBuffer, output.objectPointerBuffer, output.memoryFeaturesBuffer
}

// Without a frame command buffer, the tracker creates and commits its own:
// tracker.encodeInitialFrame(inputBuffer:prompts:) / tracker.encodeNextFrame(inputBuffer:)

tracker.reset()   // start a new sequence

Notes:

• The tracker keeps the conditioning frame plus the six most recent frames as spatial memory, and the conditioning frame plus up to fifteen recent object pointers, matching the official track_step selection.
• Output buffers are private GPU memory, written by the command buffer the frame was encoded onto. Work encoded after the tracker on that buffer, or committed after it on the same queue, reads them with no wait; blit them to a shared buffer if you need CPU access.
• If you encode a frame and then abandon the uncommitted command buffer, call reset(): the memory bank references buffers that were never written.
• A dropped frame (nil) consumes no frame index, so temporal positions stay correct across drops.
• Memory attention is one compiled graph sized for the maximum (seven memories, sixteen pointers), with an additive key mask removing unused slots from attention, so it stays a single graph however many frames are tracked. prewarmMemoryAttention() compiles it up front; otherwise it compiles on the first encodeNextFrame.
• Spatial memories are kept in float32. The official video predictor stores them as bfloat16; this package deliberately does not.
• Releasing the tracker or any stage while work is in flight is safe.

Not yet supported: correction prompts on later frames, multiple conditioning frames, direct mask conditioning, reverse tracking, and multiple objects.

Command-buffer ownership

Every stage's encode(..., commandBuffer:) and submit(..., commandBuffer:, completion:) take the caller's own MPSCommandBuffer and never wrap, commit, or wait on it. The caller commits it (Fabric's per-frame buffer is one).

MPSGraphExecutable.encode may commitAndContinue, which commits the underlying Metal buffer and swaps a new one into the same MPSCommandBuffer. Stages encode onto the instance you pass, so any number of stages, and your own blits and compute passes, can be encoded back to back onto one buffer, and you commit that buffer once yourself. Completion handlers are registered on the live underlying buffer after encoding, so they fire when the last segment completes.

run(...) is the synchronous path: it creates, commits, and waits on its own command buffer.

Accuracy tests

swift test

The differential tests compare all 262,144 image-embedding values, all 262,144 dense mask-prompt values, complete multimask logits and IoU predictions, and a non-square resized-logit output to official PyTorch outputs. The video tests run the whole tracker over eleven real frames (including a fully occluded frame and the recovery after it, and every attention shape from one memory to seven memories with ten pointers) and compare each frame's mask, IoU, object score, pointer and new memory to the official track_step, at roughly 1e-5 error. They also test an entirely GPU-resident mask-prompt → decoder → postprocessor chain with no intervening CPU waits. Fixture generators live in Tools; Tools/export_weights.py recreates the native mapped-weight files from the official checkpoint.

Roadmap

1. Correction prompts on later frames and multiple conditioning frames (needs the decoder's single-mask token mode for multi-point frames).
2. Direct mask conditioning, reverse tracking, and multi-object state with an optional non-overlap stage.

Fabric node integration should wait until the tracker's buffer-lifetime and backpressure contract has been exercised in the package.

Upstream

EfficientTAM is from yformer/EfficientTAM and is licensed under Apache 2.0. This package includes converted weights from the official efficienttam_ti_512x512.pt checkpoint.