ComfyUI/comfy/ldm/sam3/tracker.py

# SAM3 video tracker: memory encoder, memory attention, SAM mask decoder/prompt encoder.

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from tqdm import tqdm

try:
    import cv2
    _HAS_CV2 = True
except ImportError:
    from scipy import ndimage
    _HAS_CV2 = False

import comfy.model_management
from comfy.ldm.modules.attention import optimized_attention
from comfy.ldm.sam3.sam import rope_2d, PositionEmbeddingSine
from comfy.ops import cast_to_input
from comfy.ldm.flux.math import apply_rope1
from comfy.ldm.cascade.common import LayerNorm2d_op
from comfy.ldm.sam3.sam import MLP, PositionEmbeddingRandom
from comfy.ldm.sam3.sam import TwoWayTransformer as SAMTwoWayTransformer

NO_OBJ_SCORE = -1024.0


def to_spatial(x, H, W):
    """Reshape (B, H*W, C) → (B, C, H, W)."""
    return x.view(x.shape[0], H, W, -1).permute(0, 3, 1, 2)

class MultiplexState:
    """Tracks object-to-slot assignments for multiplex tracking. Provides mux/demux operations."""

    def __init__(self, num_objects, multiplex_count, device, dtype):
        self.multiplex_count = multiplex_count
        self.device = device
        self.dtype = dtype
        self._build(num_objects)

    def mux(self, x):
        """[N_obj, ...] -> [num_buckets, multiplex_count, ...]"""
        out_shape = (self.num_buckets, self.multiplex_count) + x.shape[1:]
        return (self.mux_matrix.to(device=x.device, dtype=x.dtype) @ x.reshape(self.total_valid_entries, -1)).view(out_shape)

    def demux(self, x):
        """[num_buckets, multiplex_count, ...] -> [N_obj, ...]"""
        out_shape = (self.total_valid_entries,) + x.shape[2:]
        flat = x.reshape(self.num_buckets * self.multiplex_count, -1)
        return (self.demux_matrix.to(device=x.device, dtype=x.dtype) @ flat).view(out_shape)

    def get_valid_object_mask(self):
        """[num_buckets, multiplex_count] bool tensor, True for valid slots."""
        return (self.mux_matrix.sum(dim=1) > 0).reshape(self.num_buckets, self.multiplex_count)

    def _build(self, num_objects):
        M = self.multiplex_count
        self.num_buckets = (num_objects + M - 1) // M
        self.total_valid_entries = num_objects
        total_slots = self.num_buckets * M
        self.mux_matrix = torch.zeros(total_slots, num_objects, device=self.device, dtype=self.dtype)
        self.demux_matrix = torch.zeros(num_objects, total_slots, device=self.device, dtype=self.dtype)
        oids = torch.arange(num_objects, device=self.device)
        slots = (oids // M) * M + (oids % M)
        self.mux_matrix[slots, oids] = 1.0
        self.demux_matrix[oids, slots] = 1.0

    def add_objects(self, n_new):
        """Grow multiplex state for n_new additional objects."""
        self._build(self.total_valid_entries + n_new)

def _compute_mask_overlap(masks_a, masks_b):
    """Max of IoU and IoM (intersection over minimum area). More robust to size differences."""
    a_flat = (masks_a > 0).float().flatten(1)
    b_flat = (masks_b > 0).float().flatten(1)
    intersection = a_flat @ b_flat.T
    area_a = a_flat.sum(1, keepdim=True)
    area_b = b_flat.sum(1, keepdim=True).T
    iou = intersection / (area_a + area_b - intersection).clamp(min=1)
    iom = intersection / torch.min(area_a.expand_as(iou), area_b.expand_as(iou)).clamp(min=1)
    return torch.max(iou, iom)


def _nms_masks(masks, scores, thresh=0.5):
    """Mask-based NMS using IoU+IoM overlap. Returns (filtered_masks, filtered_scores)."""
    order = scores.argsort(descending=True)
    masks, scores = masks[order], scores[order]
    keep = []
    for i in range(masks.shape[0]):
        if keep:
            if _compute_mask_overlap(masks[i:i+1], masks[torch.tensor(keep, device=masks.device)]).max() >= thresh:
                continue
        keep.append(i)
    return masks[keep], scores[keep]


def _get_connected_components(mask_bin):
    """Get connected component labels and areas. mask_bin: [B, 1, H, W] uint8."""
    labels_list, areas_list = [], []
    for i in range(mask_bin.shape[0]):
        m = mask_bin[i, 0].cpu().numpy()
        if _HAS_CV2:
            _, labeled, stats, _ = cv2.connectedComponentsWithStats(m, connectivity=8)
            areas = stats[labeled, cv2.CC_STAT_AREA].astype('int32')
        else:
            labeled, num_features = ndimage.label(m)
            areas = np.zeros_like(m, dtype=np.int32)
            for c in range(1, num_features + 1):
                component = labeled == c
                areas[component] = component.sum()
        labels_list.append(torch.from_numpy(labeled).to(mask_bin.device))
        areas_list.append(torch.from_numpy(areas).to(device=mask_bin.device, dtype=torch.int32))
    return torch.stack(labels_list).unsqueeze(1), torch.stack(areas_list).unsqueeze(1)


def fill_holes_in_mask_scores(mask, max_area=0):
    """Remove small foreground sprinkles and fill small background holes using connected components."""
    if max_area <= 0:
        return mask

    # Fill holes: small connected components in background → foreground
    mask_bg = (mask <= 0).to(torch.uint8)
    _, areas_bg = _get_connected_components(mask_bg)
    small_bg = mask_bg.bool() & (areas_bg <= max_area)
    mask = torch.where(small_bg, 0.1, mask)

    # Remove sprinkles: small connected components in foreground → background
    # Only remove if area < min(max_area, half of total foreground area)
    mask_fg = (mask > 0).to(torch.uint8)
    fg_area_thresh = mask_fg.sum(dim=(2, 3), keepdim=True, dtype=torch.int32)
    fg_area_thresh.floor_divide_(2).clamp_(max=max_area)
    _, areas_fg = _get_connected_components(mask_fg)
    small_fg = mask_fg.bool() & (areas_fg <= fg_area_thresh)
    mask = torch.where(small_fg, -0.1, mask)

    return mask


def apply_rope_memory(q, k, freqs, num_heads, num_k_exclude_rope=0):
    """Apply 2D axial RoPE to memory attention using flux rope format.

    Args:
        q: [B, Nq, C] projected queries (current frame features)
        k: [B, Nk, C] projected keys (memory tokens)
        freqs: [1, Nq, dim//2, 2, 2] flux-format rotation matrices for one frame
        num_heads: number of attention heads
        num_k_exclude_rope: number of trailing k tokens to skip RoPE (object pointers)
    """
    B, Nq, C = q.shape
    head_dim = C // num_heads

    # freqs shape: [1, 1, Nq, dim//2, 2, 2] (heads broadcast dim already included)
    q_h = q.view(B, Nq, num_heads, head_dim).transpose(1, 2)
    q_h = apply_rope1(q_h, freqs)
    q = q_h.transpose(1, 2).reshape(B, Nq, C)

    # Apply RoPE to k (excluding last num_k_exclude_rope tokens)
    Nk = k.shape[1]
    num_k_rope = Nk - num_k_exclude_rope
    if num_k_rope > 0:
        # Repeat freqs for multiple frames of spatial memory
        Nf = freqs.shape[2]  # spatial positions in one frame
        if num_k_rope > Nf:
            r = (num_k_rope + Nf - 1) // Nf
            pe_k = freqs.repeat(1, 1, r, 1, 1, 1)[:, :, :num_k_rope]
        else:
            pe_k = freqs[:, :, :num_k_rope]

        k_h = k[:, :num_k_rope].view(B, num_k_rope, num_heads, head_dim).transpose(1, 2)
        k_h = apply_rope1(k_h, pe_k)
        k = k.clone()
        k[:, :num_k_rope] = k_h.transpose(1, 2).reshape(B, num_k_rope, C)

    return q, k


def get_1d_sine_pe(pos_inds, dim, temperature=10000):
    """1D sinusoidal positional encoding for temporal positions."""
    pe_dim = dim // 2
    dim_t = torch.arange(pe_dim, dtype=torch.float32, device=pos_inds.device)
    dim_t = temperature ** (2 * (dim_t // 2) / pe_dim)
    pos_embed = pos_inds.unsqueeze(-1) / dim_t
    return torch.cat([pos_embed.sin(), pos_embed.cos()], dim=-1)


def _pad_to_buckets(tensor, target_buckets):
    """Pad a [num_buckets, ...] tensor to target_buckets along dim 0 if needed."""
    if tensor.shape[0] >= target_buckets:
        return tensor
    pad_shape = (target_buckets - tensor.shape[0],) + tensor.shape[1:]
    return torch.cat([tensor, torch.zeros(pad_shape, device=tensor.device, dtype=tensor.dtype)], dim=0)


def pack_masks(masks):
    """Pack binary masks [*, H, W] to bit-packed [*, H, W//8] uint8. W must be divisible by 8."""
    binary = masks > 0
    shifts = torch.arange(8, device=masks.device)
    return (binary.view(*masks.shape[:-1], -1, 8) * (1 << shifts)).sum(-1).byte()


def unpack_masks(packed):
    """Unpack bit-packed [*, H, W//8] uint8 to bool [*, H, W*8]."""
    bits = torch.tensor([1, 2, 4, 8, 16, 32, 64, 128], dtype=torch.uint8, device=packed.device)
    return (packed.unsqueeze(-1) & bits).bool().view(*packed.shape[:-1], -1)


def _prep_frame(images, idx, device, dt, size):
    """Slice CPU full-res frames, transfer to GPU in target dtype, and resize to (size, size)."""
    return comfy.utils.common_upscale(images[idx].to(device=device, dtype=dt), size, size, "bicubic", crop="disabled")


def _compute_backbone(backbone_fn, frame, frame_idx=None):
    """Compute backbone features for a single frame. Returns (vision_feats, vision_pos, feat_sizes, features, trunk_out)."""
    features, positions, trunk_out = backbone_fn(frame, frame_idx=frame_idx)
    feat_sizes = [(x.shape[-2], x.shape[-1]) for x in features]
    vision_feats = [x.flatten(2).permute(0, 2, 1) for x in features]
    vision_pos = [x.flatten(2).permute(0, 2, 1) for x in positions]
    return vision_feats, vision_pos, feat_sizes, features, trunk_out


def collect_memory_tokens(output_dict, frame_idx, num_maskmem, maskmem_tpos_enc, device,
                          collect_image_feats=False, tpos_v2=False, num_buckets=None):
    """Collect spatial memory, position encodings, and optionally image features from past frames."""
    to_cat_memory, to_cat_memory_pos = [], []
    to_cat_image_feat, to_cat_image_pos = [], []

    def _append(out, tpos_idx):
        feats = out["maskmem_features"].to(device)
        if num_buckets is not None:
            feats = _pad_to_buckets(feats, num_buckets)
        to_cat_memory.append(feats.flatten(2).permute(0, 2, 1))
        enc = out["maskmem_pos_enc"][-1].to(device).flatten(2).permute(0, 2, 1)
        if num_buckets is not None:
            enc = _pad_to_buckets(enc, num_buckets)
        tpos = cast_to_input(maskmem_tpos_enc[tpos_idx], enc)
        to_cat_memory_pos.append(enc + tpos)
        if collect_image_feats and "image_features" in out:
            to_cat_image_feat.append(out["image_features"].to(device))
            to_cat_image_pos.append(out["image_pos_enc"].to(device) + tpos)

    cond_outputs = output_dict["cond_frame_outputs"]
    for t, out in cond_outputs.items():
        if tpos_v2:
            t_pos = frame_idx - t
            tpos_idx = num_maskmem - t_pos - 1 if 0 < t_pos < num_maskmem else num_maskmem - 1
        else:
            tpos_idx = num_maskmem - 1
        _append(out, tpos_idx)

    for t_pos in range(1, num_maskmem):
        out = output_dict["non_cond_frame_outputs"].get(frame_idx - (num_maskmem - t_pos), None)
        if out is None or out.get("maskmem_features") is None:
            continue
        _append(out, num_maskmem - t_pos - 1)

    return to_cat_memory, to_cat_memory_pos, to_cat_image_feat, to_cat_image_pos, cond_outputs


def compute_tpos_enc(rel_pos_list, device, d_model, proj_layer, dtype=None, max_abs_pos=None):
    """Temporal position encoding for object pointers."""
    pos_enc = torch.tensor(rel_pos_list, dtype=torch.float32, device=device) / max((max_abs_pos or 2) - 1, 1)
    pos_enc = get_1d_sine_pe(pos_enc, dim=d_model)
    if dtype is not None:
        pos_enc = pos_enc.to(dtype)
    return proj_layer(pos_enc)


def forward_sam_heads(backbone_features, prompt_encoder, mask_decoder, obj_ptr_proj, no_obj_fn,
                      image_size, point_inputs=None, mask_inputs=None, box_inputs=None,
                      high_res_features=None, multimask_output=False):
    """Shared SAM prompt encoder + mask decoder forward for both SAM3 and SAM3.1 trackers."""
    device = backbone_features.device
    # Batch size from inputs (mask_inputs may have N_obj > 1 while backbone is batch 1)
    if mask_inputs is not None:
        B = mask_inputs.shape[0]
    elif box_inputs is not None:
        B = box_inputs.shape[0]
    elif point_inputs is not None:
        B = point_inputs["point_coords"].shape[0]
    else:
        B = backbone_features.shape[0]

    if point_inputs is not None:
        sam_point_coords = point_inputs["point_coords"]
        sam_point_labels = point_inputs["point_labels"]
    else:
        sam_point_coords = torch.zeros(B, 1, 2, device=device)
        sam_point_labels = -torch.ones(B, 1, dtype=torch.int32, device=device)

    if mask_inputs is not None:
        prompt_size = (prompt_encoder.image_embedding_size[0] * 4, prompt_encoder.image_embedding_size[1] * 4)
        if mask_inputs.shape[-2:] != prompt_size:
            sam_mask_prompt = F.interpolate(mask_inputs, size=prompt_size, mode="bilinear", align_corners=False, antialias=True)
        else:
            sam_mask_prompt = mask_inputs
    else:
        sam_mask_prompt = None

    sparse, dense = prompt_encoder(points=(sam_point_coords, sam_point_labels), boxes=box_inputs, masks=sam_mask_prompt)
    sparse = cast_to_input(sparse, backbone_features)
    dense = cast_to_input(dense, backbone_features)
    image_pe = cast_to_input(prompt_encoder.get_dense_pe(), backbone_features)

    low_res_multimasks, ious, sam_output_tokens, object_score_logits = mask_decoder(
        image_embeddings=backbone_features, image_pe=image_pe,
        sparse_prompt_embeddings=sparse, dense_prompt_embeddings=dense,
        high_res_features=high_res_features, multimask_output=multimask_output, return_all=True,
    )

    is_obj_appearing = object_score_logits > 0
    low_res_multimasks = torch.where(is_obj_appearing[:, None, None], low_res_multimasks,
                                      torch.tensor(NO_OBJ_SCORE, device=device, dtype=low_res_multimasks.dtype))
    high_res_multimasks = F.interpolate(low_res_multimasks, size=(image_size, image_size), mode="bilinear", align_corners=False)

    sam_output_token = sam_output_tokens[:, 0]
    if multimask_output:
        best_iou_inds = torch.argmax(ious, dim=-1)
        batch_inds = torch.arange(B, device=device)
        low_res_masks = low_res_multimasks[batch_inds, best_iou_inds].unsqueeze(1)
        high_res_masks = high_res_multimasks[batch_inds, best_iou_inds].unsqueeze(1)
        if sam_output_tokens.size(1) > 1:
            sam_output_token = sam_output_tokens[batch_inds, best_iou_inds]
    else:
        low_res_masks, high_res_masks = low_res_multimasks, high_res_multimasks

    obj_ptr = obj_ptr_proj(sam_output_token)
    obj_ptr = no_obj_fn(obj_ptr, is_obj_appearing)

    return low_res_masks, high_res_masks, obj_ptr, object_score_logits


def use_mask_as_output(backbone_features, high_res_features, mask_inputs, mask_downsample,
                       prompt_encoder, mask_decoder, obj_ptr_proj, no_obj_fn, image_size, backbone_stride):
    """Shared mask-as-output for both SAM3 and SAM3.1 trackers."""
    out_scale, out_bias = 20.0, -10.0
    mask_inputs_float = cast_to_input(mask_inputs, backbone_features)
    high_res_masks = mask_inputs_float * out_scale + out_bias
    low_res_masks = F.interpolate(high_res_masks, size=(image_size // backbone_stride * 4,) * 2,
                                   mode="bilinear", align_corners=False, antialias=True)
    _, _, obj_ptr, _ = forward_sam_heads(
        backbone_features, prompt_encoder, mask_decoder, obj_ptr_proj, no_obj_fn,
        image_size, mask_inputs=mask_downsample(mask_inputs_float), high_res_features=high_res_features,
    )
    is_obj_appearing = torch.any(mask_inputs.flatten(1) > 0.0, dim=1)[..., None]
    alpha = is_obj_appearing.to(obj_ptr.dtype)
    object_score_logits = out_scale * alpha + out_bias
    return low_res_masks, high_res_masks, obj_ptr, object_score_logits


# Split attention with configurable input dims (for asymmetric cross-attention)
class SplitAttn(nn.Module):
    def __init__(self, embed_dim, num_heads=1, kv_dim=None, internal_dim=None, device=None, dtype=None, operations=None):
        super().__init__()
        self.num_heads = num_heads
        kv_dim = kv_dim or embed_dim
        internal_dim = internal_dim or embed_dim
        self.q_proj = operations.Linear(embed_dim, internal_dim, device=device, dtype=dtype)
        self.k_proj = operations.Linear(kv_dim, internal_dim, device=device, dtype=dtype)
        self.v_proj = operations.Linear(kv_dim, internal_dim, device=device, dtype=dtype)
        self.out_proj = operations.Linear(internal_dim, embed_dim, device=device, dtype=dtype)

    def forward(self, q, k=None, v=None, rope=None, num_k_exclude_rope=0):
        if k is None:
            k = q
        if v is None:
            v = k
        q = self.q_proj(q)
        k = self.k_proj(k)
        v = self.v_proj(v)
        if rope is not None:
            q, k = apply_rope_memory(q, k, rope, self.num_heads, num_k_exclude_rope)
        out = optimized_attention(q, k, v, self.num_heads, low_precision_attention=False)
        return self.out_proj(out)


class MemoryAttnLayer(nn.Module):
    def __init__(self, d_model=256, num_heads=1, kv_dim=64, dim_ff=2048, device=None, dtype=None, operations=None):
        super().__init__()
        self.num_heads = num_heads
        self.self_attn = SplitAttn(d_model, num_heads, device=device, dtype=dtype, operations=operations)
        self.cross_attn_image = SplitAttn(d_model, num_heads, kv_dim=kv_dim, device=device, dtype=dtype, operations=operations)
        self.linear1 = operations.Linear(d_model, dim_ff, device=device, dtype=dtype)
        self.linear2 = operations.Linear(dim_ff, d_model, device=device, dtype=dtype)
        self.norm1 = operations.LayerNorm(d_model, device=device, dtype=dtype)
        self.norm2 = operations.LayerNorm(d_model, device=device, dtype=dtype)
        self.norm3 = operations.LayerNorm(d_model, device=device, dtype=dtype)

    def forward(self, x, memory, memory_pos=None, rope=None, num_k_exclude_rope=0):
        x = x + self.self_attn(self.norm1(x), rope=rope)
        mem_k = memory + memory_pos if memory_pos is not None else memory
        x = x + self.cross_attn_image(self.norm2(x), mem_k, memory, rope=rope, num_k_exclude_rope=num_k_exclude_rope)
        normed = self.norm3(x)
        x = x + self.linear2(F.relu(self.linear1(normed)))
        return x


class MemoryAttnEncoder(nn.Module):
    def __init__(self, d_model=256, num_heads=1, kv_dim=64, dim_ff=2048, num_layers=4, image_size=1008, patch_size=14,
                 device=None, dtype=None, operations=None):
        super().__init__()
        self.layers = nn.ModuleList([
            MemoryAttnLayer(d_model, num_heads, kv_dim, dim_ff, device=device, dtype=dtype, operations=operations)
            for _ in range(num_layers)
        ])
        self.norm = operations.LayerNorm(d_model, device=device, dtype=dtype)
        hw = image_size // patch_size
        self.register_buffer("_rope", rope_2d(hw, hw, d_model // num_heads), persistent=False)

    def forward(self, x, memory, src_pos=None, memory_pos=None, num_k_exclude_rope=0):
        if src_pos is not None:
            x = x + 0.1 * src_pos

        rope = self._rope.to(device=x.device)
        for layer in self.layers:
            x = layer(x, memory, memory_pos=memory_pos, rope=rope, num_k_exclude_rope=num_k_exclude_rope)
        return self.norm(x)


class MemoryTransformer(nn.Module):
    def __init__(self, d_model=256, num_heads=1, kv_dim=64, dim_ff=2048, num_layers=4, device=None, dtype=None, operations=None):
        super().__init__()
        self.encoder = MemoryAttnEncoder(d_model, num_heads, kv_dim, dim_ff, num_layers, device=device, dtype=dtype, operations=operations)


def _upscale_masks(output_upscaling, conv_s0, conv_s1, src_out, high_res_features):
    """Shared upscaling for SAM mask decoders: deconv + high-res feature integration."""
    dc1, ln1, act1, dc2, act2 = output_upscaling
    if high_res_features is not None:
        upscaled = act1(ln1(dc1(src_out) + conv_s1(high_res_features[1])))
        upscaled = act2(dc2(upscaled) + conv_s0(high_res_features[0]))
    else:
        upscaled = act2(dc2(act1(ln1(dc1(src_out)))))
    return upscaled


class SAMMaskDecoder(nn.Module):
    def __init__(self, d_model=256, num_multimask_outputs=3, device=None, dtype=None, operations=None):
        super().__init__()
        self.num_mask_tokens = num_multimask_outputs + 1

        self.transformer = SAMTwoWayTransformer(depth=2, embedding_dim=d_model, num_heads=8, mlp_dim=2048, device=device, dtype=dtype, operations=operations)

        self.iou_token = operations.Embedding(1, d_model, device=device, dtype=dtype)
        self.mask_tokens = operations.Embedding(self.num_mask_tokens, d_model, device=device, dtype=dtype)
        self.obj_score_token = operations.Embedding(1, d_model, device=device, dtype=dtype)

        # Output upscaling: d_model -> d_model//4 -> d_model//8 at 4x resolution
        LN2d = LayerNorm2d_op(operations)
        self.output_upscaling = nn.Sequential(
            operations.ConvTranspose2d(d_model, d_model // 4, kernel_size=2, stride=2, device=device, dtype=dtype), LN2d(d_model // 4, device=device, dtype=dtype), nn.GELU(),
            operations.ConvTranspose2d(d_model // 4, d_model // 8, kernel_size=2, stride=2, device=device, dtype=dtype), nn.GELU(),
        )

        # High-res feature integration
        self.conv_s0 = operations.Conv2d(d_model, d_model // 8, kernel_size=1, device=device, dtype=dtype)
        self.conv_s1 = operations.Conv2d(d_model, d_model // 4, kernel_size=1, device=device, dtype=dtype)

        # Per-mask hypernetwork MLPs
        self.output_hypernetworks_mlps = nn.ModuleList([
            MLP(d_model, d_model, d_model // 8, 3, device=device, dtype=dtype, operations=operations)
            for _ in range(self.num_mask_tokens)
        ])

        self.iou_prediction_head = MLP(d_model, d_model, self.num_mask_tokens, 3, device=device, dtype=dtype, operations=operations)
        self.pred_obj_score_head = MLP(d_model, d_model, 1, 3, device=device, dtype=dtype, operations=operations)

    def forward(self, image_embeddings, image_pe, sparse_prompt_embeddings, dense_prompt_embeddings,
                high_res_features=None, multimask_output=False, return_all=False):
        B = sparse_prompt_embeddings.shape[0]
        ref = sparse_prompt_embeddings
        # Token order: [obj_score(1), iou(1), mask(num_mask_tokens)]
        tokens = torch.cat([cast_to_input(self.obj_score_token.weight, ref),
                            cast_to_input(self.iou_token.weight, ref),
                            cast_to_input(self.mask_tokens.weight, ref)], dim=0)
        tokens = torch.cat([tokens.unsqueeze(0).expand(B, -1, -1), sparse_prompt_embeddings], dim=1)

        src = image_embeddings
        if src.shape[0] != B:
            src = src.expand(B, -1, -1, -1)
        src = src + dense_prompt_embeddings
        pos_src = image_pe.expand(B, -1, -1, -1)

        b, c, h, w = src.shape
        src_flat = src.flatten(2).permute(0, 2, 1)
        pos_flat = pos_src.flatten(2).permute(0, 2, 1)

        hs, src_out = self.transformer(src_flat, pos_flat, tokens)

        obj_score_token_out = hs[:, 0, :]
        iou_token_out = hs[:, 1, :]
        mask_tokens_out = hs[:, 2:2 + self.num_mask_tokens, :]

        src_out = src_out.permute(0, 2, 1).view(b, c, h, w)
        upscaled = _upscale_masks(self.output_upscaling, self.conv_s0, self.conv_s1, src_out, high_res_features)

        hyper_in = torch.stack([
            mlp(mask_tokens_out[:, i, :]) for i, mlp in enumerate(self.output_hypernetworks_mlps)
        ], dim=1)

        masks = (hyper_in @ upscaled.flatten(2)).view(B, self.num_mask_tokens, upscaled.shape[2], upscaled.shape[3])
        iou_pred = self.iou_prediction_head(iou_token_out)
        object_score_logits = self.pred_obj_score_head(obj_score_token_out)

        if multimask_output:
            out_masks = masks[:, 1:]
            out_iou = iou_pred[:, 1:]
            out_tokens = mask_tokens_out[:, 1:]
        else:
            out_masks = masks[:, 0:1]
            out_iou = iou_pred[:, 0:1]
            out_tokens = mask_tokens_out[:, 0:1]

        if return_all:
            return out_masks, out_iou, out_tokens, object_score_logits
        return out_masks, out_iou


class SAMPromptEncoder(nn.Module):
    def __init__(self, d_model=256, image_embedding_size=(72, 72), input_image_size=(1008, 1008), device=None, dtype=None, operations=None):
        super().__init__()
        self.embed_dim = d_model
        self.image_embedding_size = image_embedding_size
        self.input_image_size = input_image_size

        self.pe_layer = PositionEmbeddingRandom(d_model // 2)
        self.point_embeddings = nn.ModuleList([
            operations.Embedding(1, d_model, device=device, dtype=dtype) for _ in range(4)
        ])
        self.not_a_point_embed = operations.Embedding(1, d_model, device=device, dtype=dtype)

        LN2d = LayerNorm2d_op(operations)
        self.mask_downscaling = nn.Sequential(
            operations.Conv2d(1, 4, kernel_size=2, stride=2, device=device, dtype=dtype),
            LN2d(4, device=device, dtype=dtype), nn.GELU(),
            operations.Conv2d(4, 16, kernel_size=2, stride=2, device=device, dtype=dtype),
            LN2d(16, device=device, dtype=dtype), nn.GELU(),
            operations.Conv2d(16, d_model, kernel_size=1, device=device, dtype=dtype),
        )
        self.no_mask_embed = operations.Embedding(1, d_model, device=device, dtype=dtype)

    def get_dense_pe(self):
        return self.pe_layer(self.image_embedding_size)

    def forward(self, points=None, boxes=None, masks=None):
        ref = points[0] if points is not None else boxes if boxes is not None else masks
        B = 1
        sparse = torch.empty((B, 0, self.embed_dim), device=ref.device, dtype=ref.dtype)

        if points is not None:
            coords, labels = points
            B = coords.shape[0]
            # Pad with an extra point (label=-1) when no boxes are provided (matching reference)
            if boxes is None:
                coords = torch.cat([coords, torch.zeros(B, 1, 2, device=coords.device, dtype=coords.dtype)], dim=1)
                labels = torch.cat([labels, -torch.ones(B, 1, device=labels.device, dtype=labels.dtype)], dim=1)
            pe = self.pe_layer.forward_with_coords(coords + 0.5, self.input_image_size)
            for i in range(4):
                pe[labels == i] += cast_to_input(self.point_embeddings[i].weight, ref)
            invalid = (labels == -1)
            pe[invalid] = 0.0
            pe[invalid] += cast_to_input(self.not_a_point_embed.weight, ref)
            sparse = torch.cat([sparse.expand(B, -1, -1), pe], dim=1)

        if boxes is not None:
            B = boxes.shape[0]
            corners = self.pe_layer.forward_with_coords((boxes.reshape(-1, 2, 2) + 0.5), self.input_image_size)
            corners[:, 0] += cast_to_input(self.point_embeddings[2].weight, ref)
            corners[:, 1] += cast_to_input(self.point_embeddings[3].weight, ref)
            sparse = torch.cat([sparse.expand(B, -1, -1), corners], dim=1)

        if masks is not None:
            dense = self.mask_downscaling(masks)
        else:
            dense = cast_to_input(self.no_mask_embed.weight, ref).reshape(1, -1, 1, 1).expand(
                B, -1, self.image_embedding_size[0], self.image_embedding_size[1])

        return sparse, dense


class CXBlock(nn.Module):
    def __init__(self, dim=256, kernel_size=7, device=None, dtype=None, operations=None):
        super().__init__()
        self.dwconv = operations.Conv2d(dim, dim, kernel_size=kernel_size, padding=kernel_size // 2, groups=dim, device=device, dtype=dtype)
        self.norm = operations.LayerNorm(dim, device=device, dtype=dtype)
        self.pwconv1 = operations.Linear(dim, 4 * dim, device=device, dtype=dtype)
        self.pwconv2 = operations.Linear(4 * dim, dim, device=device, dtype=dtype)
        self.gamma = nn.Parameter(torch.ones(dim, device=device, dtype=dtype))

    def forward(self, x):
        residual = x
        x = self.dwconv(x).permute(0, 2, 3, 1)
        x = self.pwconv2(F.gelu(self.pwconv1(self.norm(x))))
        x.mul_(cast_to_input(self.gamma, x))
        return residual + x.permute(0, 3, 1, 2)


class MaskDownSampler(nn.Module):
    def __init__(self, out_dim=256, in_chans=1, channels=None, interpol_size=(1152, 1152), device=None, dtype=None, operations=None):
        super().__init__()
        self.interpol_size = list(interpol_size) if interpol_size else None
        if channels is None:
            channels = [4, 16, 64, out_dim]  # SAM3 default
        LN2d = LayerNorm2d_op(operations)
        layers = []
        prev = in_chans
        for ch in channels:
            layers += [operations.Conv2d(prev, ch, kernel_size=3, stride=2, padding=1, device=device, dtype=dtype),
                       LN2d(ch, device=device, dtype=dtype), nn.GELU()]
            prev = ch
        layers.append(operations.Conv2d(prev, out_dim, kernel_size=1, device=device, dtype=dtype))
        self.encoder = nn.Sequential(*layers)

    def forward(self, x):
        if self.interpol_size is not None and list(x.shape[-2:]) != self.interpol_size:
            x = F.interpolate(x, size=self.interpol_size, mode="bilinear", align_corners=False, antialias=True)
        return self.encoder(x)


class Fuser(nn.Module):
    def __init__(self, dim=256, num_layers=2, device=None, dtype=None, operations=None):
        super().__init__()
        self.layers = nn.Sequential(*[CXBlock(dim, device=device, dtype=dtype, operations=operations) for _ in range(num_layers)])

    def forward(self, x):
        return self.layers(x)


# --- SAM3.1 Multiplex components ---

class DecoupledMemoryAttnLayer(nn.Module):
    """Decoupled cross-attention layer for SAM3.1: fuses image and memory projections."""

    def __init__(self, d_model=256, num_heads=1, dim_ff=2048, device=None, dtype=None, operations=None):
        super().__init__()
        self.num_heads = num_heads
        # Self-attention projections (flat, not nested)
        self.self_attn_q_proj = operations.Linear(d_model, d_model, device=device, dtype=dtype)
        self.self_attn_k_proj = operations.Linear(d_model, d_model, device=device, dtype=dtype)
        self.self_attn_v_proj = operations.Linear(d_model, d_model, device=device, dtype=dtype)
        self.self_attn_out_proj = operations.Linear(d_model, d_model, device=device, dtype=dtype)
        # Cross-attention projections
        self.cross_attn_q_proj = operations.Linear(d_model, d_model, device=device, dtype=dtype)
        self.cross_attn_k_proj = operations.Linear(d_model, d_model, device=device, dtype=dtype)
        self.cross_attn_v_proj = operations.Linear(d_model, d_model, device=device, dtype=dtype)
        self.cross_attn_out_proj = operations.Linear(d_model, d_model, device=device, dtype=dtype)
        # Image cross-attention (q/k only, fused with cross_attn)
        self.image_cross_attn_q_proj = operations.Linear(d_model, d_model, device=device, dtype=dtype)
        self.image_cross_attn_k_proj = operations.Linear(d_model, d_model, device=device, dtype=dtype)
        # FFN
        self.linear1 = operations.Linear(d_model, dim_ff, device=device, dtype=dtype)
        self.linear2 = operations.Linear(dim_ff, d_model, device=device, dtype=dtype)
        self.norm1 = operations.LayerNorm(d_model, device=device, dtype=dtype)
        self.norm2 = operations.LayerNorm(d_model, device=device, dtype=dtype)
        self.norm3 = operations.LayerNorm(d_model, device=device, dtype=dtype)

    def forward(self, image, x, memory_image, memory, memory_image_pos=None,
                rope=None, num_k_exclude_rope=0):
        # Self-attention with RoPE
        normed = self.norm1(x)
        q = self.self_attn_q_proj(normed)
        k = self.self_attn_k_proj(normed)
        v = self.self_attn_v_proj(normed)
        if rope is not None:
            q, k = apply_rope_memory(q, k, rope, self.num_heads, 0)
        x = x + self.self_attn_out_proj(optimized_attention(q, k, v, self.num_heads, low_precision_attention=False))

        # Decoupled cross-attention: fuse image and memory projections
        normed = self.norm2(x)
        q = self.image_cross_attn_q_proj(image) + self.cross_attn_q_proj(normed)
        k = self.image_cross_attn_k_proj(memory_image) + self.cross_attn_k_proj(memory)
        if memory_image_pos is not None:
            k = k + memory_image_pos
        v = self.cross_attn_v_proj(memory)
        if rope is not None:
            q, k = apply_rope_memory(q, k, rope, self.num_heads, num_k_exclude_rope)
        x = x + self.cross_attn_out_proj(optimized_attention(q, k, v, self.num_heads, low_precision_attention=False))

        # FFN
        x = x + self.linear2(F.gelu(self.linear1(self.norm3(x))))
        return image, x


class DecoupledMemoryEncoder(nn.Module):
    """Memory attention encoder for SAM3.1 with decoupled cross-attention."""

    def __init__(self, d_model=256, num_heads=1, dim_ff=2048, num_layers=4, image_size=1008, patch_size=14,
                 device=None, dtype=None, operations=None):
        super().__init__()
        self.layers = nn.ModuleList([
            DecoupledMemoryAttnLayer(d_model, num_heads, dim_ff, device=device, dtype=dtype, operations=operations)
            for _ in range(num_layers)
        ])
        self.norm = operations.LayerNorm(d_model, device=device, dtype=dtype)
        hw = image_size // patch_size
        self.register_buffer("_rope", rope_2d(hw, hw, d_model // num_heads), persistent=False)

    def forward(self, x, memory, memory_pos=None, src_pos=None, num_k_exclude_rope=0,
                memory_image=None, memory_image_pos=None):
        image = x  # constant residual for decoupled cross-attention
        output = x
        if src_pos is not None:
            output = output + 0.1 * src_pos

        B, _, C = x.shape
        rope = self._rope.to(device=x.device)

        # memory_image: raw backbone features from past frames for decoupled cross-attention
        if memory_image is None:
            # Fallback: use spatial portion of memory (without obj pointers)
            num_spatial = memory.shape[1] - num_k_exclude_rope
            memory_image = memory[:, :num_spatial]
            memory_image_pos = memory_pos[:, :num_spatial] if memory_pos is not None else None
        # Pad memory_image to match memory length (zeros for obj pointer tokens)
        if memory_image.shape[1] < memory.shape[1]:
            pad_len = memory.shape[1] - memory_image.shape[1]
            pad = torch.zeros(B, pad_len, C, device=memory.device, dtype=memory.dtype)
            memory_image = torch.cat([memory_image, pad], dim=1)
            if memory_image_pos is not None:
                ptr_pos = memory_pos[:, -pad_len:] if memory_pos is not None else torch.zeros_like(pad)
                memory_image_pos = torch.cat([memory_image_pos, ptr_pos], dim=1)

        for layer in self.layers:
            image, output = layer(image, output, memory_image, memory,
                                  memory_image_pos=memory_image_pos, rope=rope,
                                  num_k_exclude_rope=num_k_exclude_rope)

        return self.norm(output)


class DecoupledMemoryTransformer(nn.Module):
    def __init__(self, d_model=256, num_heads=1, dim_ff=2048, num_layers=4, device=None, dtype=None, operations=None):
        super().__init__()
        self.encoder = DecoupledMemoryEncoder(d_model, num_heads, dim_ff, num_layers,
                                              device=device, dtype=dtype, operations=operations)


class MemoryBackbone(nn.Module):
    """Memory encoder: downsamples mask, fuses with pixel features, optionally compresses."""

    def __init__(self, d_model=256, out_dim=None, in_chans=1, channels=None, device=None, dtype=None, operations=None):
        super().__init__()
        self.mask_downsampler = MaskDownSampler(d_model, in_chans=in_chans, channels=channels, device=device, dtype=dtype, operations=operations)
        self.pix_feat_proj = operations.Conv2d(d_model, d_model, kernel_size=1, device=device, dtype=dtype)
        self.fuser = Fuser(d_model, num_layers=2, device=device, dtype=dtype, operations=operations)
        self.has_out_proj = out_dim is not None and out_dim != d_model
        if self.has_out_proj:
            self.out_proj = operations.Conv2d(d_model, out_dim, kernel_size=1, device=device, dtype=dtype)
            feat_dim = out_dim
        else:
            feat_dim = d_model
        self.position_encoding = PositionEmbeddingSine(num_pos_feats=feat_dim, normalize=True)

    def forward(self, image_features, mask_for_mem, skip_mask_sigmoid=False):
        if not skip_mask_sigmoid:
            mask_for_mem = mask_for_mem.sigmoid()
        mask_features = self.mask_downsampler(cast_to_input(mask_for_mem, image_features))
        if mask_features.shape[-2:] != image_features.shape[-2:]:
            mask_features = F.interpolate(mask_features, size=image_features.shape[-2:], mode="bilinear", align_corners=False)
        features = self.pix_feat_proj(image_features) + mask_features
        features = self.fuser(features)
        if self.has_out_proj:
            features = self.out_proj(features)
        pos = cast_to_input(self.position_encoding(features), features)
        return {"vision_features": features, "vision_pos_enc": [pos]}


class MultiplexMaskDecoder(nn.Module):
    """SAM mask decoder for SAM3.1 multiplex: predicts masks for num_multiplex objects simultaneously.

    Uses multimask_outputs_only=True: num_mask_output_per_object = num_multimask_outputs (no +1).
    Hypernetwork MLPs are shared across multiplex objects.
    Token order: [obj_score_token(M), iou_token(M), mask_tokens(M*T)].
    """

    def __init__(self, d_model=256, num_multiplex=16, num_multimask_outputs=3, device=None, dtype=None, operations=None):
        super().__init__()
        self.num_multiplex = num_multiplex
        self.num_mask_output_per_object = num_multimask_outputs  # 3 (multimask_outputs_only)
        total_mask_tokens = num_multiplex * self.num_mask_output_per_object  # 48

        self.transformer = SAMTwoWayTransformer(depth=2, embedding_dim=d_model, num_heads=8, mlp_dim=2048, device=device, dtype=dtype, operations=operations)

        self.obj_score_token = operations.Embedding(num_multiplex, d_model, device=device, dtype=dtype)
        self.iou_token = operations.Embedding(num_multiplex, d_model, device=device, dtype=dtype)
        self.mask_tokens = operations.Embedding(total_mask_tokens, d_model, device=device, dtype=dtype)

        LN2d = LayerNorm2d_op(operations)
        self.output_upscaling = nn.Sequential(
            operations.ConvTranspose2d(d_model, d_model // 4, kernel_size=2, stride=2, device=device, dtype=dtype),
            LN2d(d_model // 4, device=device, dtype=dtype), nn.GELU(),
            operations.ConvTranspose2d(d_model // 4, d_model // 8, kernel_size=2, stride=2, device=device, dtype=dtype), nn.GELU(),
        )
        self.conv_s0 = operations.Conv2d(d_model, d_model // 8, kernel_size=1, device=device, dtype=dtype)
        self.conv_s1 = operations.Conv2d(d_model, d_model // 4, kernel_size=1, device=device, dtype=dtype)

        # Shared across all multiplex objects (one per mask output)
        self.output_hypernetworks_mlps = nn.ModuleList([
            MLP(d_model, d_model, d_model // 8, 3, device=device, dtype=dtype, operations=operations)
            for _ in range(self.num_mask_output_per_object)
        ])
        self.iou_prediction_head = MLP(d_model, d_model, self.num_mask_output_per_object, 3, device=device, dtype=dtype, operations=operations)
        self.pred_obj_score_head = MLP(d_model, d_model, 1, 3, device=device, dtype=dtype, operations=operations)

    def forward(self, image_embeddings, image_pe, sparse_prompt_embeddings, dense_prompt_embeddings,
                high_res_features=None, multimask_output=False, return_all=False, extra_per_object_embeddings=None):
        B = sparse_prompt_embeddings.shape[0]
        M = self.num_multiplex
        T = self.num_mask_output_per_object

        # Token order: [obj_score(M), iou(M), mask(M*T)]
        ref = sparse_prompt_embeddings
        mask_tokens = cast_to_input(self.mask_tokens.weight, ref)
        if extra_per_object_embeddings is not None:
            mask_tokens = mask_tokens.view(1, M, T, -1).expand(B, -1, -1, -1) + extra_per_object_embeddings.unsqueeze(2)
            mask_tokens = mask_tokens.flatten(1, 2)  # [B, M*T, C]
            other_tokens = torch.cat([cast_to_input(self.obj_score_token.weight, ref),
                                      cast_to_input(self.iou_token.weight, ref)], dim=0).unsqueeze(0).expand(B, -1, -1)
            tokens = torch.cat([other_tokens, mask_tokens, sparse_prompt_embeddings], dim=1)
        else:
            tokens = torch.cat([cast_to_input(self.obj_score_token.weight, ref),
                                cast_to_input(self.iou_token.weight, ref), mask_tokens], dim=0)
            tokens = torch.cat([tokens.unsqueeze(0).expand(B, -1, -1), sparse_prompt_embeddings], dim=1)

        src = image_embeddings
        if src.shape[0] != B:
            src = src.expand(B, -1, -1, -1)
        src = src + dense_prompt_embeddings
        pos_src = image_pe.expand(B, -1, -1, -1)

        b, c, h, w = src.shape
        hs, src_out = self.transformer(src.flatten(2).permute(0, 2, 1), pos_src.flatten(2).permute(0, 2, 1), tokens)

        # Parse output tokens
        obj_score_token_out = hs[:, :M]
        iou_token_out = hs[:, M:2 * M]
        mask_tokens_out = hs[:, 2 * M:2 * M + M * T]

        src_out = src_out.permute(0, 2, 1).view(b, c, h, w)
        upscaled = _upscale_masks(self.output_upscaling, self.conv_s0, self.conv_s1, src_out, high_res_features)

        # Reshape mask tokens to [B, M, T, C] and apply shared hypernetwork MLPs per mask output index
        mask_tokens_2d = mask_tokens_out.view(B, M, T, -1)
        hyper_in = torch.stack([
            self.output_hypernetworks_mlps[i](mask_tokens_2d[:, :, i, :])  # [B, M, C//8]
            for i in range(T)
        ], dim=2)  # [B, M, T, C//8]

        # Generate masks: [B, M*T, H*W] -> [B, M, T, H, W]
        masks = torch.bmm(hyper_in.flatten(1, 2), upscaled.flatten(2)).view(b, M, T, upscaled.shape[2], upscaled.shape[3])

        # IoU and object scores
        iou_pred = self.iou_prediction_head(iou_token_out).view(b, M, T)
        object_score_logits = self.pred_obj_score_head(obj_score_token_out)  # [B, M, 1]

        # multimask_outputs_only: always output all T masks (no singlemask token)
        sam_tokens_out = mask_tokens_2d[:, :, 0:1]  # [B, M, 1, C]

        if return_all:
            return masks, iou_pred, sam_tokens_out, object_score_logits
        return masks, iou_pred


class SAM3Tracker(nn.Module):
    def __init__(self, d_model=256, mem_dim=64, num_maskmem=7, device=None, dtype=None, operations=None, **kwargs):
        super().__init__()

        # Memory attention transformer
        self.transformer = MemoryTransformer(d_model, num_heads=1, kv_dim=mem_dim, dim_ff=2048, num_layers=4,
                                             device=device, dtype=dtype, operations=operations)
        # SAM components
        self.sam_mask_decoder = SAMMaskDecoder(d_model, device=device, dtype=dtype, operations=operations)
        self.sam_prompt_encoder = SAMPromptEncoder(d_model, device=device, dtype=dtype, operations=operations)

        # Memory backbone
        self.maskmem_backbone = MemoryBackbone(d_model, out_dim=mem_dim, device=device, dtype=dtype, operations=operations)

        # Standalone parameters
        self.maskmem_tpos_enc = nn.Parameter(torch.zeros(num_maskmem, 1, 1, mem_dim, device=device, dtype=dtype))
        self.no_mem_embed = nn.Parameter(torch.zeros(1, 1, d_model, device=device, dtype=dtype))
        self.register_buffer("no_mem_pos_enc", torch.zeros(1, 1, d_model, device=device, dtype=dtype))  # checkpoint key, unused in forward
        self.no_obj_embed_spatial = nn.Parameter(torch.zeros(1, mem_dim, device=device, dtype=dtype))
        self.no_obj_ptr = nn.Parameter(torch.zeros(1, d_model, device=device, dtype=dtype))

        # Object pointer projection
        self.obj_ptr_proj = MLP(d_model, d_model, d_model, 3, device=device, dtype=dtype, operations=operations)
        self.obj_ptr_tpos_proj = operations.Linear(d_model, mem_dim, device=device, dtype=dtype)

        # Mask downsample: Conv2d stride 4 to reduce GT mask to SAM logit scale
        self.mask_downsample = operations.Conv2d(1, 1, kernel_size=4, stride=4, device=device, dtype=dtype)

        # Config
        self.d_model = d_model
        self.mem_dim = mem_dim
        self.num_maskmem = num_maskmem
        self.image_size = 1008
        self.backbone_stride = 14
        self.max_obj_ptrs_in_encoder = 16
        self.sigmoid_scale_for_mem_enc = 20.0
        self.sigmoid_bias_for_mem_enc = -10.0

    def _no_obj_blend(self, obj_ptr, is_obj):
        alpha = is_obj.to(obj_ptr.dtype)
        return torch.lerp(cast_to_input(self.no_obj_ptr, obj_ptr), obj_ptr, alpha)

    def _forward_sam_heads(self, backbone_features, point_inputs=None, mask_inputs=None, box_inputs=None,
                           high_res_features=None, multimask_output=False):
        return forward_sam_heads(backbone_features, self.sam_prompt_encoder, self.sam_mask_decoder,
                                 self.obj_ptr_proj, self._no_obj_blend, self.image_size,
                                 point_inputs, mask_inputs, box_inputs, high_res_features, multimask_output)

    def _use_mask_as_output(self, backbone_features, high_res_features, mask_inputs):
        return use_mask_as_output(backbone_features, high_res_features, mask_inputs,
                                  self.mask_downsample, self.sam_prompt_encoder, self.sam_mask_decoder,
                                  self.obj_ptr_proj, self._no_obj_blend, self.image_size, self.backbone_stride)

    def _prepare_memory_conditioned_features(self, frame_idx, is_init_cond_frame, current_vision_feats, current_vision_pos_embeds, feat_sizes, output_dict, num_frames):
        """Fuse current frame features with memory from previous frames."""
        B = current_vision_feats[-1].shape[0]
        C = self.d_model
        H, W = feat_sizes[-1]
        device = current_vision_feats[-1].device

        if self.num_maskmem == 0:
            return current_vision_feats[-1].permute(0, 2, 1).view(B, C, H, W)

        if is_init_cond_frame:
            # First conditioning frame: no memory yet, add no_mem_embed
            pix_feat = current_vision_feats[-1] + cast_to_input(self.no_mem_embed, current_vision_feats[-1])
            return to_spatial(pix_feat, H, W)

        to_cat_memory, to_cat_memory_pos, _, _, cond_outputs = collect_memory_tokens(
            output_dict, frame_idx, self.num_maskmem, self.maskmem_tpos_enc, device)

        max_obj_ptrs = min(num_frames, self.max_obj_ptrs_in_encoder)
        pos_and_ptrs = []
        for t, out in cond_outputs.items():
            if t <= frame_idx:
                pos_and_ptrs.append(((frame_idx - t), out["obj_ptr"].to(device)))
        for t_diff in range(1, max_obj_ptrs):
            t = frame_idx - t_diff
            if t < 0:
                break
            out = output_dict["non_cond_frame_outputs"].get(t, None)
            if out is not None:
                pos_and_ptrs.append((t_diff, out["obj_ptr"].to(device)))

        num_obj_ptr_tokens = 0
        if len(pos_and_ptrs) > 0:
            pos_list, ptrs_list = zip(*pos_and_ptrs)
            obj_ptrs = torch.stack(ptrs_list, dim=1)  # [B, N, C=256]

            # Temporal position encoding for pointers
            obj_pos = compute_tpos_enc(
                list(pos_list), device, self.d_model, self.obj_ptr_tpos_proj,
                max_abs_pos=max_obj_ptrs, dtype=current_vision_feats[-1].dtype
            )  # [N, mem_dim=64]
            obj_pos = obj_pos.unsqueeze(0).expand(B, -1, -1)  # [B, N, 64]

            # Split each 256-dim pointer into 4 x 64-dim tokens
            if self.mem_dim < C:
                N = obj_ptrs.shape[1]
                obj_ptrs = obj_ptrs.view(B, N, C // self.mem_dim, self.mem_dim)  # [B, N, 4, 64]
                obj_ptrs = obj_ptrs.reshape(B, N * (C // self.mem_dim), self.mem_dim)  # [B, N*4, 64]
                obj_pos = obj_pos.unsqueeze(2).expand(-1, -1, C // self.mem_dim, -1)
                obj_pos = obj_pos.reshape(B, N * (C // self.mem_dim), self.mem_dim)  # [B, N*4, 64]

            to_cat_memory.append(obj_ptrs)
            to_cat_memory_pos.append(obj_pos)
            num_obj_ptr_tokens = obj_ptrs.shape[1]

        if len(to_cat_memory) == 0:
            # No memory available yet, add no_mem_embed
            pix_feat = current_vision_feats[-1] + cast_to_input(self.no_mem_embed, current_vision_feats[-1])
            return to_spatial(pix_feat, H, W)

        # Concatenate all memory and position encodings [B, total_mem, mem_dim=64]
        memory = torch.cat(to_cat_memory, dim=1)
        memory_pos = torch.cat(to_cat_memory_pos, dim=1)

        # Run memory attention encoder
        pix_feat = current_vision_feats[-1]  # [B, HW, C]
        src_pos = current_vision_pos_embeds[-1]  # [B, HW, C]

        pix_feat_with_mem = self.transformer.encoder(
            x=pix_feat,
            memory=memory,
            src_pos=src_pos,
            memory_pos=memory_pos,
            num_k_exclude_rope=num_obj_ptr_tokens,
        )
        return to_spatial(pix_feat_with_mem, H, W)

    def _encode_new_memory(self, pix_feat, pred_masks_high_res, object_score_logits, is_mask_from_pts=False):
        """Encode predicted mask into memory features."""
        if is_mask_from_pts:
            mask_for_mem = (pred_masks_high_res > 0).to(pix_feat.dtype)
        else:
            mask_for_mem = torch.sigmoid(pred_masks_high_res)

        mask_for_mem.mul_(self.sigmoid_scale_for_mem_enc).add_(self.sigmoid_bias_for_mem_enc)

        maskmem_out = self.maskmem_backbone(pix_feat, mask_for_mem, skip_mask_sigmoid=True)
        maskmem_features = maskmem_out["vision_features"]
        maskmem_pos_enc = maskmem_out["vision_pos_enc"]

        # Add no_obj_embed for occluded objects
        alpha = (object_score_logits > 0).to(maskmem_features.dtype)[..., None, None]
        no_obj = cast_to_input(self.no_obj_embed_spatial, maskmem_features)[..., None, None].expand_as(maskmem_features)
        return maskmem_features + (1 - alpha) * no_obj, maskmem_pos_enc

    def track_step(self, frame_idx, is_init_cond_frame, current_vision_feats, current_vision_pos_embeds, feat_sizes, mask_inputs, output_dict,
                   num_frames, point_inputs=None):
        """Track one frame: fuse with memory, predict mask, encode memory."""
        current_out = {}

        # High-res features for SAM head [stride-8, stride-4]
        if len(current_vision_feats) > 1:
            high_res_features = [
                x.view(x.shape[0], feat_sizes[i][0], feat_sizes[i][1], -1).permute(0, 3, 1, 2)
                for i, x in enumerate(current_vision_feats[:-1])
            ]
        else:
            high_res_features = None

        # Top-level feature for memory
        H, W = feat_sizes[-1]

        if mask_inputs is not None:
            # Conditioning frame: use mask directly
            pix_feat = to_spatial(current_vision_feats[-1], H, W)
            sam_outputs = self._use_mask_as_output(pix_feat, high_res_features, mask_inputs)
        else:
            # Track frame: fuse with memory, then SAM decoder
            pix_feat_with_mem = self._prepare_memory_conditioned_features(
                frame_idx=frame_idx,
                is_init_cond_frame=is_init_cond_frame,
                current_vision_feats=current_vision_feats,
                current_vision_pos_embeds=current_vision_pos_embeds,
                feat_sizes=feat_sizes,
                output_dict=output_dict,
                num_frames=num_frames,
            )
            # Use multimask for point prompts on init frames (picks best of 3 candidates)
            num_pts = 0 if point_inputs is None else point_inputs["point_labels"].size(1)
            multimask_output = is_init_cond_frame and 0 < num_pts <= 1
            sam_outputs = self._forward_sam_heads(
                backbone_features=pix_feat_with_mem,
                point_inputs=point_inputs,
                high_res_features=high_res_features,
                multimask_output=multimask_output,
            )

        (low_res_masks, high_res_masks, obj_ptr, object_score_logits) = sam_outputs

        # Clean low-res masks: remove sprinkles and fill holes
        low_res_masks = fill_holes_in_mask_scores(low_res_masks, max_area=200)
        high_res_masks = F.interpolate(low_res_masks, size=(self.image_size, self.image_size), mode="bilinear", align_corners=False)

        current_out["pred_masks"] = low_res_masks
        current_out["pred_masks_high_res"] = high_res_masks
        current_out["obj_ptr"] = obj_ptr
        current_out["object_score_logits"] = object_score_logits

        # Encode memory
        if self.num_maskmem > 0:
            pix_feat = to_spatial(current_vision_feats[-1], H, W)
            maskmem_features, maskmem_pos_enc = self._encode_new_memory(
                pix_feat=pix_feat,
                pred_masks_high_res=high_res_masks,
                object_score_logits=object_score_logits,
                is_mask_from_pts=(point_inputs is not None),
            )
            current_out["maskmem_features"] = maskmem_features
            current_out["maskmem_pos_enc"] = maskmem_pos_enc
        else:
            current_out["maskmem_features"] = None
            current_out["maskmem_pos_enc"] = None

        return current_out

    def _compute_backbone_frame(self, backbone_fn, frame, frame_idx=None):
        vision_feats, vision_pos, feat_sizes, _, _ = _compute_backbone(backbone_fn, frame, frame_idx)
        # SAM3: drop last FPN level
        return vision_feats[:-1], vision_pos[:-1], feat_sizes[:-1]

    def _track_single_object(self, backbone_fn, images, initial_mask, pbar=None,
                             target_device=None, target_dtype=None):
        """Track one object, computing backbone per frame to save VRAM."""
        N = images.shape[0]
        device = target_device if target_device is not None else images.device
        dt = target_dtype if target_dtype is not None else images.dtype
        size = self.image_size
        output_dict = {"cond_frame_outputs": {}, "non_cond_frame_outputs": {}}
        all_masks = []

        for frame_idx in tqdm(range(N), desc="tracking"):
            vision_feats, vision_pos, feat_sizes = self._compute_backbone_frame(
                backbone_fn, _prep_frame(images, slice(frame_idx, frame_idx + 1), device, dt, size), frame_idx=frame_idx)
            mask_input = None
            if frame_idx == 0:
                mask_input = F.interpolate(initial_mask.to(device=device, dtype=dt),
                    size=(self.image_size, self.image_size), mode="bilinear", align_corners=False)
                mask_input = (mask_input > 0.5).to(dt)

            current_out = self.track_step(
                frame_idx=frame_idx, is_init_cond_frame=(frame_idx == 0),
                current_vision_feats=vision_feats, current_vision_pos_embeds=vision_pos,
                feat_sizes=feat_sizes, mask_inputs=mask_input, output_dict=output_dict, num_frames=N)

            if frame_idx == 0:
                output_dict["cond_frame_outputs"][frame_idx] = current_out
            else:
                output_dict["non_cond_frame_outputs"][frame_idx] = current_out
                lookback = max(self.num_maskmem, self.max_obj_ptrs_in_encoder)
                for old_idx in list(output_dict["non_cond_frame_outputs"]):
                    if old_idx < frame_idx - lookback:
                        del output_dict["non_cond_frame_outputs"][old_idx]
            # Move masks to CPU immediately to free VRAM
            all_masks.append(current_out["pred_masks_high_res"].to(comfy.model_management.intermediate_device()))
            if pbar is not None:
                pbar.update(1)

        return torch.cat(all_masks, dim=0)  # [N, 1, H, W]

    def track_video(self, backbone_fn, images, initial_masks, pbar=None,
                    target_device=None, target_dtype=None, **kwargs):
        """Track one or more objects across video frames.

        Args:
            backbone_fn: callable that returns (sam2_features, sam2_positions, trunk_out) for a frame
            images: [N, 3, H, W] CPU full-res video frames (resized per-frame to self.image_size)
            initial_masks: [N_obj, 1, H, W] binary masks for first frame (one per object)
            pbar: optional progress bar

        Returns:
            [N, N_obj, image_size, image_size] predicted mask logits per frame per object
        """
        N_obj = initial_masks.shape[0]
        per_object = []
        for obj_idx in range(N_obj):
            obj_masks = self._track_single_object(
                backbone_fn, images, initial_masks[obj_idx:obj_idx + 1], pbar=pbar,
                target_device=target_device, target_dtype=target_dtype)
            per_object.append(obj_masks)

        return torch.cat(per_object, dim=1)  # [N, N_obj, H, W]


class SAM31Tracker(nn.Module):
    """SAM3.1 multiplex tracker: decoupled memory attention, dual decoder, 16-object multiplex."""

    def __init__(self, d_model=256, mem_dim=256, num_maskmem=7, num_multiplex=16, device=None, dtype=None, operations=None, **kwargs):
        super().__init__()
        self.d_model = d_model
        self.mem_dim = mem_dim
        self.num_maskmem = num_maskmem
        self.num_multiplex = num_multiplex
        self.image_size = 1008
        self.backbone_stride = 14
        self.max_obj_ptrs_in_encoder = 16
        self.sigmoid_scale_for_mem_enc = 2.0
        self.sigmoid_bias_for_mem_enc = -1.0

        # Memory attention (decoupled cross-attention, 8 heads matching reference)
        self.transformer = DecoupledMemoryTransformer(d_model, num_heads=8, dim_ff=2048, num_layers=4,
                                                      device=device, dtype=dtype, operations=operations)

        # Propagation decoder (multiplex: 16 objects, multimask_outputs_only)
        self.sam_mask_decoder = MultiplexMaskDecoder(d_model, num_multiplex, num_multimask_outputs=3,
                                                     device=device, dtype=dtype, operations=operations)
        # Interactive decoder (single object, same as SAM3)
        self.interactive_sam_mask_decoder = SAMMaskDecoder(d_model, num_multimask_outputs=3,
                                                           device=device, dtype=dtype, operations=operations)
        self.interactive_sam_prompt_encoder = SAMPromptEncoder(d_model, device=device, dtype=dtype, operations=operations)

        # Memory backbone (mem_dim=256, no out_proj compression)
        self.maskmem_backbone = MemoryBackbone(d_model, in_chans=num_multiplex * 2, channels=[16, 64, 256, 1024],
                                                device=device, dtype=dtype, operations=operations)

        # Standalone parameters
        self.maskmem_tpos_enc = nn.Parameter(torch.zeros(num_maskmem, 1, 1, mem_dim, device=device, dtype=dtype))
        self.no_obj_embed_spatial = nn.Parameter(torch.zeros(num_multiplex, mem_dim, device=device, dtype=dtype))
        self.interactivity_no_mem_embed = nn.Parameter(torch.zeros(1, 1, d_model, device=device, dtype=dtype))

        # Object pointer projection
        self.obj_ptr_proj = MLP(d_model, d_model, d_model, 3, device=device, dtype=dtype, operations=operations)
        self.obj_ptr_tpos_proj = operations.Linear(d_model, mem_dim, device=device, dtype=dtype)
        self.no_obj_ptr_linear = operations.Linear(d_model, d_model, device=device, dtype=dtype)
        self.interactive_obj_ptr_proj = MLP(d_model, d_model, d_model, 3, device=device, dtype=dtype, operations=operations)

        # Interactive mask downsample
        self.interactive_mask_downsample = operations.Conv2d(1, 1, kernel_size=4, stride=4, device=device, dtype=dtype)

        # Multiplex validity embeddings
        self.output_valid_embed = nn.Parameter(torch.zeros(num_multiplex, d_model, device=device, dtype=dtype))
        self.output_invalid_embed = nn.Parameter(torch.zeros(num_multiplex, d_model, device=device, dtype=dtype))

        # Position encoding for image (used by multiplex decoder)
        self.image_pe_layer = PositionEmbeddingRandom(d_model // 2)

    def _no_obj_blend(self, obj_ptr, is_obj):
        alpha = is_obj.to(obj_ptr.dtype)
        return torch.lerp(self.no_obj_ptr_linear(obj_ptr), obj_ptr, alpha)

    def _forward_sam_heads(self, backbone_features, point_inputs=None, mask_inputs=None, box_inputs=None,
                           high_res_features=None, multimask_output=False):
        return forward_sam_heads(backbone_features, self.interactive_sam_prompt_encoder, self.interactive_sam_mask_decoder,
                                 self.interactive_obj_ptr_proj, self._no_obj_blend, self.image_size,
                                 point_inputs, mask_inputs, box_inputs, high_res_features, multimask_output)

    def _use_mask_as_output(self, backbone_features, high_res_features, mask_inputs):
        return use_mask_as_output(backbone_features, high_res_features, mask_inputs,
                                  self.interactive_mask_downsample, self.interactive_sam_prompt_encoder,
                                  self.interactive_sam_mask_decoder, self.interactive_obj_ptr_proj,
                                  self._no_obj_blend, self.image_size, self.backbone_stride)

    def _prepare_memory_conditioned_features(self, frame_idx, is_init_cond_frame, current_vision_feats,
                                              current_vision_pos_embeds, feat_sizes, output_dict, num_frames,
                                              multiplex_state=None):
        B = current_vision_feats[-1].shape[0]
        C = self.d_model
        H, W = feat_sizes[-1]
        device = current_vision_feats[-1].device
        num_buc = multiplex_state.num_buckets if multiplex_state is not None else None

        if self.num_maskmem == 0:
            return current_vision_feats[-1].permute(0, 2, 1).view(B, C, H, W)

        if is_init_cond_frame:
            pix_feat = current_vision_feats[-1] + cast_to_input(self.interactivity_no_mem_embed, current_vision_feats[-1])
            return to_spatial(pix_feat, H, W)

        to_cat_memory, to_cat_memory_pos, to_cat_image_feat, to_cat_image_pos, cond_outputs = collect_memory_tokens(
            output_dict, frame_idx, self.num_maskmem, self.maskmem_tpos_enc, device,
            collect_image_feats=True, tpos_v2=True, num_buckets=num_buc)

        max_obj_ptrs = min(num_frames, self.max_obj_ptrs_in_encoder)
        pos_and_ptrs = []
        for t, out in cond_outputs.items():
            if t <= frame_idx and "obj_ptr" in out:
                ptr = out["obj_ptr"].to(device)
                if num_buc is not None:
                    ptr = _pad_to_buckets(ptr, num_buc)
                pos_and_ptrs.append(((frame_idx - t), ptr))
        for t_diff in range(1, max_obj_ptrs):
            t = frame_idx - t_diff
            if t < 0:
                break
            out = output_dict["non_cond_frame_outputs"].get(t, None)
            if out is not None and "obj_ptr" in out:
                ptr = out["obj_ptr"].to(device)
                if num_buc is not None:
                    ptr = _pad_to_buckets(ptr, num_buc)
                pos_and_ptrs.append((t_diff, ptr))

        num_obj_ptr_tokens = 0
        if len(pos_and_ptrs) > 0:
            pos_list, ptrs_list = zip(*pos_and_ptrs)
            obj_ptrs = torch.stack(ptrs_list, dim=1)  # [num_buckets, N, M, C]
            B_ptr = obj_ptrs.shape[0]
            N_ptrs = obj_ptrs.shape[1]
            M = obj_ptrs.shape[2]
            obj_ptrs = obj_ptrs.reshape(B_ptr, N_ptrs * M, -1)
            obj_pos = compute_tpos_enc(list(pos_list), device, self.d_model, self.obj_ptr_tpos_proj,
                                       max_abs_pos=max_obj_ptrs, dtype=current_vision_feats[-1].dtype)
            obj_pos = obj_pos.unsqueeze(0).expand(B_ptr, -1, -1)
            obj_pos = obj_pos.unsqueeze(2).expand(-1, -1, M, -1).reshape(B_ptr, N_ptrs * M, -1)
            to_cat_memory.append(obj_ptrs)
            to_cat_memory_pos.append(obj_pos)
            num_obj_ptr_tokens = obj_ptrs.shape[1]

        if len(to_cat_memory) == 0:
            pix_feat = current_vision_feats[-1] + cast_to_input(self.interactivity_no_mem_embed, current_vision_feats[-1])
            return to_spatial(pix_feat, H, W)

        memory = torch.cat(to_cat_memory, dim=1)
        memory_pos = torch.cat(to_cat_memory_pos, dim=1)

        # Expand vision features to num_buckets if memory has more buckets than B
        mem_B = memory.shape[0]
        x = current_vision_feats[-1]
        x_pos = current_vision_pos_embeds[-1]
        if x.shape[0] < mem_B:
            x = x.expand(mem_B, -1, -1)
            x_pos = x_pos.expand(mem_B, -1, -1)

        if len(to_cat_image_feat) > 0:
            # Decoupled cross-attention: separate image features from memory
            memory_image = cast_to_input(torch.cat(to_cat_image_feat, dim=1), x)
            memory_image_pos = cast_to_input(torch.cat(to_cat_image_pos, dim=1), x)
            if memory_image.shape[0] < mem_B:
                memory_image = memory_image.expand(mem_B, -1, -1)
                memory_image_pos = memory_image_pos.expand(mem_B, -1, -1)
            pix_feat_with_mem = self.transformer.encoder(
                x=x,
                memory=cast_to_input(memory, x),
                memory_pos=cast_to_input(memory_pos, x),
                src_pos=cast_to_input(x_pos, x),
                num_k_exclude_rope=num_obj_ptr_tokens,
                memory_image=memory_image,
                memory_image_pos=memory_image_pos,
            )
        else:
            pix_feat_with_mem = self.transformer.encoder(
                x=x,
                memory=memory,
                memory_pos=memory_pos,
                src_pos=x_pos,
                num_k_exclude_rope=num_obj_ptr_tokens,
            )
        return to_spatial(pix_feat_with_mem, H, W)

    def _encode_new_memory(self, pix_feat, pred_masks_high_res, object_score_logits, is_mask_from_pts=False,
                           multiplex_state=None, is_conditioning=False, cond_obj_mask=None):
        if is_mask_from_pts:
            mask_for_mem = (pred_masks_high_res > 0).to(pix_feat.dtype)
        else:
            mask_for_mem = torch.sigmoid(pred_masks_high_res)
        mask_for_mem.mul_(self.sigmoid_scale_for_mem_enc).add_(self.sigmoid_bias_for_mem_enc)

        # Mux masks: [N_obj, 1, H, W] -> [num_buckets, M, H, W]
        mux_masks = multiplex_state.mux(mask_for_mem[:, 0])

        # Conditioning channel: 1.0 = clean mask (trust it), 0.0 = propagation (noisy)
        N_obj = mask_for_mem.shape[0]
        cond_values = torch.full((N_obj,), 0.0, device=mask_for_mem.device, dtype=mask_for_mem.dtype)
        if is_conditioning:
            cond_values[:] = 1.0
        elif cond_obj_mask is not None:
            cond_values[cond_obj_mask] = 1.0
        cond_spatial = cond_values.view(-1, 1, 1, 1).expand_as(mask_for_mem[:, 0:1, :, :]).squeeze(1)
        mux_cond = multiplex_state.mux(cond_spatial)  # [num_buckets, M, H, W]
        mux_input = torch.cat([mux_masks, mux_cond], dim=1)  # [num_buckets, 2*M, H, W]

        maskmem_out = self.maskmem_backbone(pix_feat, mux_input, skip_mask_sigmoid=True)
        maskmem_features = maskmem_out["vision_features"]
        maskmem_pos_enc = maskmem_out["vision_pos_enc"]

        # Add no_obj_embed_spatial for occluded objects
        is_obj = (object_score_logits > 0).float()  # [N_obj, 1]
        mux_is_obj = multiplex_state.mux(is_obj)  # [num_buckets, M, 1]
        no_obj_embed = cast_to_input(self.no_obj_embed_spatial, maskmem_features)  # [M, C]
        no_obj_spatial = no_obj_embed.unsqueeze(0)[..., None, None]  # [1, M, C, 1, 1]
        # Expand and sum across multiplex slots weighted by (1 - is_obj)
        alpha = mux_is_obj[..., None, None]  # [num_buckets, M, 1, 1, 1]
        per_slot_no_obj = ((1 - alpha) * no_obj_spatial).sum(dim=1)  # [num_buckets, C, 1, 1]
        maskmem_features = maskmem_features + per_slot_no_obj.expand_as(maskmem_features)

        return maskmem_features, maskmem_pos_enc

    def _forward_propagation(self, backbone_features, high_res_features=None, multiplex_state=None):
        """Propagation path using the multiplex SAM decoder (no prompts)."""
        B = backbone_features.shape[0]
        device = backbone_features.device

        # Suppression embeddings from valid object mask
        valid_mask = cast_to_input(multiplex_state.get_valid_object_mask().unsqueeze(-1).float(), backbone_features)
        output_valid = cast_to_input(self.output_valid_embed, backbone_features).unsqueeze(0)
        output_invalid = cast_to_input(self.output_invalid_embed, backbone_features).unsqueeze(0)
        extra_embed = valid_mask * output_valid + (1 - valid_mask) * output_invalid

        image_pe = self.image_pe_layer((backbone_features.shape[-2], backbone_features.shape[-1]), device=backbone_features.device)
        image_pe = cast_to_input(image_pe, backbone_features)

        masks, iou_pred, sam_tokens_out, object_score_logits = self.sam_mask_decoder(
            image_embeddings=backbone_features, image_pe=image_pe,
            sparse_prompt_embeddings=torch.empty(B, 0, self.d_model, device=device, dtype=backbone_features.dtype),
            dense_prompt_embeddings=torch.zeros(B, self.d_model, *backbone_features.shape[-2:], device=device, dtype=backbone_features.dtype),
            high_res_features=high_res_features, multimask_output=True, return_all=True,
            extra_per_object_embeddings=extra_embed.expand(B, -1, -1),
        )
        # masks: [B=num_buckets, M, T, H, W]
        # Demux to per-object: [N_obj, T, H, W]
        masks_obj = multiplex_state.demux(masks)
        iou_obj = multiplex_state.demux(iou_pred)
        score_obj = multiplex_state.demux(object_score_logits)
        tokens_obj = multiplex_state.demux(sam_tokens_out)

        # Select best mask by IoU for each object
        best_idx = torch.argmax(iou_obj, dim=-1)  # [N_obj]
        N_obj = masks_obj.shape[0]
        obj_range = torch.arange(N_obj, device=device)
        low_res_masks = masks_obj[obj_range, best_idx].unsqueeze(1)  # [N_obj, 1, H, W]
        # Suppress masks for objects with low confidence
        is_obj = score_obj > 0
        low_res_masks = torch.where(is_obj[:, :, None, None], low_res_masks,
                                     torch.tensor(NO_OBJ_SCORE, device=device, dtype=low_res_masks.dtype))
        high_res_masks = F.interpolate(low_res_masks.float(), size=(self.image_size, self.image_size), mode="bilinear", align_corners=False)

        # Object pointer: compute per-object, mux for storage as [num_buckets, M, C]
        sam_token = tokens_obj[:, 0]  # [N_obj, C]
        obj_ptr = self.obj_ptr_proj(sam_token)
        is_obj = (score_obj > 0).float()
        no_obj = self.no_obj_ptr_linear(obj_ptr)
        obj_ptr = is_obj * obj_ptr + (1 - is_obj) * no_obj
        obj_ptr_muxed = multiplex_state.mux(obj_ptr)  # [num_buckets, M, C]

        return low_res_masks, high_res_masks, obj_ptr_muxed, score_obj

    def track_step(self, frame_idx, is_init_cond_frame, current_vision_feats, current_vision_pos_embeds,
                   feat_sizes, mask_inputs, output_dict, num_frames, point_inputs=None,
                   interactive_high_res=None, interactive_backbone=None, propagation_high_res=None,
                   multiplex_state=None, run_mem_encoder=True):
        current_out = {}
        H, W = feat_sizes[-1]

        if mask_inputs is not None:
            # Conditioning frame: use interactive features if available, else propagation
            if interactive_backbone is not None:
                pix_feat = interactive_backbone
                # Add no_mem_embed for interactive path
                pix_flat = pix_feat.flatten(2)
                bf = pix_flat.permute(0, 2, 1) + cast_to_input(self.interactivity_no_mem_embed, pix_flat)
                pix_feat = to_spatial(bf, H, W)
                hi_res = interactive_high_res
            else:
                # Fallback: interactive backbone not available (e.g. called outside track_video).
                # Propagation features work but may produce lower-quality conditioning.
                pix_feat = to_spatial(current_vision_feats[-1], H, W)
                hi_res = propagation_high_res
            sam_outputs = self._use_mask_as_output(pix_feat, hi_res, mask_inputs)
        elif point_inputs is not None:
            # Interactive path: use interactive SAM decoder
            pix_feat_with_mem = self._prepare_memory_conditioned_features(
                frame_idx=frame_idx, is_init_cond_frame=is_init_cond_frame,
                current_vision_feats=current_vision_feats, current_vision_pos_embeds=current_vision_pos_embeds,
                feat_sizes=feat_sizes, output_dict=output_dict, num_frames=num_frames,
                multiplex_state=multiplex_state,
            )
            hi_res = interactive_high_res if interactive_high_res is not None else propagation_high_res
            num_pts = point_inputs["point_labels"].size(1)
            multimask_output = is_init_cond_frame and 0 < num_pts <= 1
            sam_outputs = self._forward_sam_heads(
                backbone_features=pix_feat_with_mem, point_inputs=point_inputs,
                high_res_features=hi_res, multimask_output=multimask_output,
            )
        else:
            # Propagation path: use multiplex SAM decoder with propagation features
            pix_feat_with_mem = self._prepare_memory_conditioned_features(
                frame_idx=frame_idx, is_init_cond_frame=is_init_cond_frame,
                current_vision_feats=current_vision_feats, current_vision_pos_embeds=current_vision_pos_embeds,
                feat_sizes=feat_sizes, output_dict=output_dict, num_frames=num_frames,
                multiplex_state=multiplex_state,
            )
            sam_outputs = self._forward_propagation(pix_feat_with_mem, propagation_high_res,
                                                     multiplex_state=multiplex_state)

        (low_res_masks, high_res_masks, obj_ptr, object_score_logits) = sam_outputs

        # Mux obj_ptr if it came from interactive path (shape [B, C]) vs propagation ([num_buckets, M, C])
        if multiplex_state is not None and obj_ptr.dim() == 2:
            obj_ptr = multiplex_state.mux(obj_ptr)  # [N_obj, C] -> [num_buckets, M, C]

        # Encode memory (can be deferred with run_mem_encoder=False)
        if run_mem_encoder and self.num_maskmem > 0:
            pix_feat = to_spatial(current_vision_feats[-1], H, W)
            maskmem_features, maskmem_pos_enc = self._encode_new_memory(
                pix_feat=pix_feat, pred_masks_high_res=high_res_masks,
                object_score_logits=object_score_logits,
                is_mask_from_pts=(point_inputs is not None),
                multiplex_state=multiplex_state,
                is_conditioning=(mask_inputs is not None),
            )
            current_out["maskmem_features"] = maskmem_features
            current_out["maskmem_pos_enc"] = maskmem_pos_enc
        else:
            current_out["maskmem_features"] = None
            current_out["maskmem_pos_enc"] = None

        # Store propagation image features for decoupled memory attention
        current_out["image_features"] = current_vision_feats[-1]  # [B, HW, C]
        current_out["image_pos_enc"] = current_vision_pos_embeds[-1]  # [B, HW, C]

        current_out["pred_masks"] = low_res_masks
        current_out["pred_masks_high_res"] = high_res_masks
        current_out["obj_ptr"] = obj_ptr
        current_out["object_score_logits"] = object_score_logits

        return current_out

    def _compute_backbone_frame(self, backbone_fn, frame, frame_idx=None):
        vision_feats, vision_pos, feat_sizes, features, trunk_out = _compute_backbone(backbone_fn, frame, frame_idx)
        return vision_feats, vision_pos, feat_sizes, list(features[:-1]), trunk_out

    @staticmethod
    def _suppress_recently_occluded(low_res_masks, last_occluded, frame_idx, threshold=0.3):
        """Suppress overlapping masks for objects that were most recently occluded.
        Prevents corrupted masks from occluded objects from contaminating other objects."""
        N_obj = low_res_masks.shape[0]
        if N_obj <= 1:
            return low_res_masks
        binary = low_res_masks[:, 0] > 0  # [N_obj, H, W]
        iou = _compute_mask_overlap(low_res_masks[:, 0], low_res_masks[:, 0])
        overlapping = torch.triu(iou >= threshold, diagonal=1)  # [N, N] upper triangle
        last_occ_i = last_occluded.unsqueeze(1)  # [N, 1]
        last_occ_j = last_occluded.unsqueeze(0)  # [1, N]
        # Suppress the more recently occluded object in each overlapping pair
        suppress_i = overlapping & (last_occ_i > last_occ_j) & (last_occ_j > -1)
        suppress_j = overlapping & (last_occ_j > last_occ_i) & (last_occ_i > -1)
        to_suppress = suppress_i.any(dim=1) | suppress_j.any(dim=0)
        # Update last_occluded for occluded/suppressed objects
        is_empty = ~binary.any(dim=(-1, -2))
        newly_occluded = is_empty | to_suppress
        last_occluded[newly_occluded] = frame_idx
        # Suppress masks
        low_res_masks[to_suppress] = -10.0
        return low_res_masks

    def _deferred_memory_encode(self, current_out, N_obj, vision_feats, feat_sizes, mux_state, device,
                                cond_obj_mask=None):
        """Deferred memory encoding for propagation frames. cond_obj_mask: per-object bool for conditioning."""
        low_res_masks = current_out["pred_masks"]  # [N_obj, 1, H_low, W_low]

        if N_obj > 1:
            lr = low_res_masks.squeeze(1)  # [N_obj, H, W]
            max_obj = torch.argmax(lr, dim=0, keepdim=True)
            batch_inds = torch.arange(N_obj, device=device)[:, None, None]
            pixel_nol = torch.where(max_obj == batch_inds, lr, torch.clamp(lr, max=-10.0))
            area_before = (lr > 0).sum(dim=(-1, -2)).float().clamp(min=1)
            area_after = (pixel_nol > 0).sum(dim=(-1, -2)).float()
            shrink_ok = (area_after / area_before) >= 0.3
            low_res_masks = torch.where(
                shrink_ok[:, None, None, None].expand_as(low_res_masks),
                low_res_masks, torch.clamp(low_res_masks, max=-10.0))

        interpol_size = self.maskmem_backbone.mask_downsampler.interpol_size
        mem_masks = F.interpolate(low_res_masks, size=interpol_size,
                                  mode="bilinear", align_corners=False)

        obj_scores = torch.where(
            (mem_masks > 0).any(dim=(-1, -2)), 10.0, -10.0)

        pix_feat = to_spatial(vision_feats[-1], feat_sizes[-1][0], feat_sizes[-1][1])
        maskmem_features, maskmem_pos_enc = self._encode_new_memory(
            pix_feat=pix_feat, pred_masks_high_res=mem_masks,
            object_score_logits=obj_scores,
            multiplex_state=mux_state, cond_obj_mask=cond_obj_mask)
        current_out["maskmem_features"] = maskmem_features
        current_out["maskmem_pos_enc"] = maskmem_pos_enc

    def _add_detected_objects(self, new_masks, mux_state, vision_feats, feat_sizes, current_out):
        """Grow MultiplexState with new detections, merge masks, re-encode memory. Modifies current_out."""
        n_old = mux_state.total_valid_entries
        mux_state.add_objects(new_masks.shape[0])
        N_obj = mux_state.total_valid_entries
        # Stored memory with old bucket counts is padded at read time by _pad_to_buckets
        for k in ("pred_masks", "pred_masks_high_res"):
            det = F.interpolate(new_masks.unsqueeze(1), size=current_out[k].shape[-2:],
                                mode="bilinear", align_corners=False)
            current_out[k] = torch.cat([current_out[k], det], dim=0)
        if self.num_maskmem > 0:
            # Mark new objects as conditioning (clean detection masks) so model trusts them
            cond_mask = torch.zeros(N_obj, dtype=torch.bool, device=new_masks.device)
            cond_mask[n_old:] = True
            self._deferred_memory_encode(current_out, N_obj, vision_feats, feat_sizes,
                                         mux_state, new_masks.device, cond_obj_mask=cond_mask)

    def _condition_with_masks(self, masks, frame_idx, vision_feats, vision_pos, feat_sizes,
                              high_res_prop, output_dict, N, mux_state, backbone_obj, frame,
                              trunk_out, threshold=0.5):
        """Condition tracker with masks on a frame."""
        mask_input = F.interpolate(masks if masks.dim() == 4 else masks.unsqueeze(1),
            size=(self.image_size, self.image_size), mode="bilinear", align_corners=False)
        mask_input = (mask_input > threshold).to(masks.dtype)
        hi_res = lo_feat = None
        if backbone_obj is not None and backbone_obj.multiplex:
            _, _, itf, _ = backbone_obj(frame, tracker_mode="interactive", cached_trunk=trunk_out, tracker_only=True)
            hi_res, lo_feat = itf[:-1], itf[-1]
        current_out = self.track_step(
            frame_idx=frame_idx, is_init_cond_frame=True, current_vision_feats=vision_feats,
            current_vision_pos_embeds=vision_pos, feat_sizes=feat_sizes, mask_inputs=mask_input,
            output_dict=output_dict, num_frames=N, interactive_high_res=hi_res,
            interactive_backbone=lo_feat, propagation_high_res=high_res_prop,
            multiplex_state=mux_state, run_mem_encoder=True)
        output_dict["cond_frame_outputs"][frame_idx] = current_out
        return current_out

    def _match_and_add_detections(self, det_masks, det_scores, current_out, mux_state,
                                  vision_feats, feat_sizes, device, max_objects=0,
                                  keep_alive=None):
        """Match detections against tracked masks, add new objects, recondition degraded tracks.
        Updates keep_alive counters: +1 for matched tracks, -1 for unmatched."""
        N_obj = mux_state.total_valid_entries
        if det_masks.shape[0] == 0:
            if keep_alive is not None:
                for i in range(N_obj):
                    keep_alive[i] = max(-4, keep_alive.get(i, 0) - 1)
            return []

        # Match at low-res (like reference)
        trk_masks = current_out["pred_masks"][:, 0]  # [N_obj, H_low, W_low]
        det_resized = F.interpolate(det_masks.unsqueeze(1), size=trk_masks.shape[-2:],
                                    mode="bilinear", align_corners=False)[:, 0]
        overlap = _compute_mask_overlap(det_resized, trk_masks)

        # Update keep_alive and find matched tracks
        matched = set()
        if overlap.shape[1] > 0:
            matched = set((overlap >= 0.5).any(dim=0).nonzero(as_tuple=True)[0].tolist())
        if keep_alive is not None:
            for i in range(N_obj):
                if i in matched:
                    keep_alive[i] = min(8, keep_alive.get(i, 0) + 1)
                else:
                    keep_alive[i] = max(-4, keep_alive.get(i, 0) - 1)

        # Recondition: high-confidence detections (>=0.8) with high overlap refresh tracked masks
        reconditioned = False
        if det_scores is not None and overlap.shape[1] > 0:
            HIGH_CONF = 0.8
            for det_idx in range(overlap.shape[0]):
                if det_scores[det_idx] < HIGH_CONF:
                    continue
                best_trk = overlap[det_idx].argmax().item()
                if overlap[det_idx, best_trk] >= 0.5:
                    # Replace tracked mask with fresh detection mask
                    current_out["pred_masks"][best_trk] = det_resized[det_idx].unsqueeze(0)
                    det_hr = F.interpolate(det_masks[det_idx:det_idx+1].unsqueeze(1),
                        size=current_out["pred_masks_high_res"].shape[-2:],
                        mode="bilinear", align_corners=False)
                    current_out["pred_masks_high_res"][best_trk] = det_hr[0]
                    reconditioned = True

        # Re-encode memory if any tracks were reconditioned
        if reconditioned and self.num_maskmem > 0:
            self._deferred_memory_encode(current_out, N_obj, vision_feats, feat_sizes, mux_state, device)

        # Add new detections (not matching any track)
        if max_objects > 0 and N_obj >= max_objects:
            return []
        max_overlap = overlap.max(dim=1)[0] if overlap.shape[1] > 0 else torch.zeros(overlap.shape[0], device=device)
        new_dets = max_overlap < 0.5
        if new_dets.any():
            if max_objects > 0:
                slots = max_objects - N_obj
                new_dets = new_dets & (torch.cumsum(new_dets.int(), 0) <= slots)
            self._add_detected_objects(det_masks[new_dets], mux_state,
                                       vision_feats, feat_sizes, current_out)
            if keep_alive is not None:
                for i in range(N_obj, mux_state.total_valid_entries):
                    keep_alive[i] = 1
            return det_scores[new_dets].tolist() if det_scores is not None else [0.0] * new_dets.sum().item()
        return []

    INTERNAL_MAX_OBJECTS = 64  # Hard ceiling on accumulated tracks; max_objects=0 or any value above this is clamped here.

    def track_video_with_detection(self, backbone_fn, images, initial_masks, detect_fn=None,
                                   new_det_thresh=0.5, max_objects=0, detect_interval=1,
                                   backbone_obj=None, pbar=None, target_device=None, target_dtype=None):
        """Track with optional per-frame detection. Returns [N, max_N_obj, H, W] mask logits."""
        if max_objects <= 0 or max_objects > self.INTERNAL_MAX_OBJECTS:
            max_objects = self.INTERNAL_MAX_OBJECTS
        N = images.shape[0]
        device = target_device if target_device is not None else images.device
        dt = target_dtype if target_dtype is not None else images.dtype
        size = self.image_size
        output_dict = {"cond_frame_outputs": {}, "non_cond_frame_outputs": {}}
        all_masks = []
        idev = comfy.model_management.intermediate_device()
        mux_state = None
        if initial_masks is not None:
            mux_state = MultiplexState(initial_masks.shape[0], self.num_multiplex, device, dt)
        obj_scores = []  # per-object detection score (1.0 for initial masks)
        keep_alive = {} if detect_fn is not None else None
        last_occluded = torch.empty(0, device=device, dtype=torch.long)  # per-object last occluded frame

        # Prefetch next frame's backbone on a separate CUDA stream
        prefetch = False
        backbone_stream = None
        if comfy.model_management.is_device_cuda(device):
            try:
                backbone_stream = torch.cuda.Stream(device=device)
                prefetch = True
            except RuntimeError:
                pass
        cur_bb = self._compute_backbone_frame(backbone_fn, _prep_frame(images, slice(0, 1), device, dt, size), frame_idx=0)

        for frame_idx in tqdm(range(N), desc="tracking"):
            vision_feats, vision_pos, feat_sizes, high_res_prop, trunk_out = cur_bb

            # Start next frame's backbone on separate stream (overlaps with current frame's work)
            if prefetch and frame_idx + 1 < N:
                backbone_stream.wait_stream(torch.cuda.current_stream(device))
                with torch.cuda.stream(backbone_stream):
                    next_bb = self._compute_backbone_frame(
                        backbone_fn, _prep_frame(images, slice(frame_idx + 1, frame_idx + 2), device, dt, size), frame_idx=frame_idx + 1)

            # Per-frame detection with NMS (skip if no detect_fn, or interval/max not met)
            det_masks = torch.empty(0, device=device)
            det_scores = None
            run_det = (detect_fn is not None
                       and frame_idx % max(detect_interval, 1) == 0
                       and not (max_objects > 0 and mux_state is not None
                                and mux_state.total_valid_entries >= max_objects))
            if run_det:
                det_out = detect_fn(trunk_out)
                scores = det_out["scores"][0].sigmoid()
                keep = scores > new_det_thresh
                det_masks, det_scores = det_out["masks"][0][keep], scores[keep]
                if det_masks.shape[0] > 1:
                    det_masks, det_scores = _nms_masks(det_masks, det_scores)

            if frame_idx == 0 and initial_masks is not None:
                current_out = self._condition_with_masks(
                    initial_masks.to(device=device, dtype=dt), frame_idx, vision_feats, vision_pos,
                    feat_sizes, high_res_prop, output_dict, N, mux_state, backbone_obj,
                    _prep_frame(images, slice(frame_idx, frame_idx + 1), device, dt, size), trunk_out)
                last_occluded = torch.full((mux_state.total_valid_entries,), -1, device=device, dtype=torch.long)
                obj_scores = [1.0] * mux_state.total_valid_entries
                if keep_alive is not None:
                    for i in range(mux_state.total_valid_entries):
                        keep_alive[i] = 8
            elif mux_state is None or mux_state.total_valid_entries == 0:
                if det_masks.shape[0] > 0:
                    if max_objects > 0:
                        det_scores = det_scores[:max_objects]
                        det_masks = det_masks[:max_objects]
                    mux_state = MultiplexState(det_masks.shape[0], self.num_multiplex, device, dt)
                    current_out = self._condition_with_masks(
                        det_masks, frame_idx, vision_feats, vision_pos, feat_sizes, high_res_prop,
                        output_dict, N, mux_state, backbone_obj,
                        _prep_frame(images, slice(frame_idx, frame_idx + 1), device, dt, size), trunk_out, threshold=0.0)
                    last_occluded = torch.full((mux_state.total_valid_entries,), -1, device=device, dtype=torch.long)
                    obj_scores = det_scores[:mux_state.total_valid_entries].tolist()
                    if keep_alive is not None:
                        for i in range(mux_state.total_valid_entries):
                            keep_alive[i] = 1
                else:
                    all_masks.append(None)
                    if pbar is not None:
                        pbar.update(1)
                    # Skip to backbone advance at end of loop
                    if frame_idx + 1 < N:
                        if prefetch:
                            torch.cuda.current_stream(device).wait_stream(backbone_stream)
                            cur_bb = next_bb
                        else:
                            cur_bb = self._compute_backbone_frame(backbone_fn, _prep_frame(images, slice(frame_idx + 1, frame_idx + 2), device, dt, size), frame_idx=frame_idx + 1)
                    continue
            else:
                N_obj = mux_state.total_valid_entries
                current_out = self.track_step(
                    frame_idx=frame_idx, is_init_cond_frame=False, current_vision_feats=vision_feats,
                    current_vision_pos_embeds=vision_pos, feat_sizes=feat_sizes, mask_inputs=None,
                    output_dict=output_dict, num_frames=N, propagation_high_res=high_res_prop,
                    multiplex_state=mux_state, run_mem_encoder=False)
                current_out["pred_masks"] = fill_holes_in_mask_scores(
                    current_out["pred_masks"], max_area=16)
                if last_occluded.shape[0] == N_obj and N_obj > 1:
                    self._suppress_recently_occluded(
                        current_out["pred_masks"], last_occluded, frame_idx)
                if self.num_maskmem > 0:
                    self._deferred_memory_encode(current_out, N_obj, vision_feats, feat_sizes, mux_state, device)
                output_dict["non_cond_frame_outputs"][frame_idx] = current_out
                lookback = max(self.num_maskmem, self.max_obj_ptrs_in_encoder)
                for old_idx in list(output_dict["non_cond_frame_outputs"]):
                    if old_idx < frame_idx - lookback:
                        del output_dict["non_cond_frame_outputs"][old_idx]
                n_before = mux_state.total_valid_entries
                new_obj_scores = self._match_and_add_detections(det_masks, det_scores, current_out, mux_state,
                                               vision_feats, feat_sizes, device, max_objects,
                                               keep_alive if run_det else None)
                n_added = mux_state.total_valid_entries - n_before
                if n_added > 0:
                    last_occluded = torch.cat([last_occluded,
                        torch.full((n_added,), -1, device=device, dtype=torch.long)])
                    obj_scores.extend(new_obj_scores)

            masks_out = current_out["pred_masks_high_res"][:, 0]
            if keep_alive is not None:
                for i in range(masks_out.shape[0]):
                    if keep_alive.get(i, 0) <= 0:
                        masks_out[i] = NO_OBJ_SCORE
            N_obj_now = mux_state.total_valid_entries if mux_state is not None else 0
            if N_obj_now > 0:
                all_masks.append(pack_masks(masks_out).to(idev))
            else:
                all_masks.append(None)
            if pbar is not None:
                pbar.update(1)

            # Next frame's backbone
            if frame_idx + 1 < N:
                if prefetch:
                    torch.cuda.current_stream(device).wait_stream(backbone_stream)
                    cur_bb = next_bb
                else:
                    cur_bb = self._compute_backbone_frame(backbone_fn, _prep_frame(images, slice(frame_idx + 1, frame_idx + 2), device, dt, size), frame_idx=frame_idx + 1)

        if not all_masks or all(m is None for m in all_masks):
            return {"packed_masks": None, "n_frames": N, "scores": []}

        max_obj = max(m.shape[0] for m in all_masks if m is not None)
        sample = next(m for m in all_masks if m is not None)
        empty_packed = torch.zeros(max_obj, *sample.shape[1:], dtype=torch.uint8, device=sample.device)
        for i, m in enumerate(all_masks):
            if m is None:
                all_masks[i] = empty_packed
            elif m.shape[0] < max_obj:
                pad = torch.zeros(max_obj - m.shape[0], *m.shape[1:], dtype=torch.uint8, device=m.device)
                all_masks[i] = torch.cat([m, pad], dim=0)
        return {"packed_masks": torch.stack(all_masks, dim=0), "n_frames": N, "scores": obj_scores}