Blame for engine.py

"""3AM inference engine for the Hugging Face Space.

3AM routing-fusion model: SAM2.1 Hiera-L whose memory attention is wrapped by
RoutingMemoryAttention (routing tokens + FiLM fusion of MUSt3R features), MUSt3R-512 streamed with
memory_mode='kv' and an explicit memory policy (two anchor frames, a write stride, and an optional
cap on the number of memory frames; by default every frame's KV is kept; no keyframe selection,
no KD-tree, no retrieval), and a SAM2 memory policy that keeps the previous frame plus the most
recent positive masks.

Frames are letterboxed (never cropped) to a MUSt3R-supported aspect ratio; masks are mapped
back to the original resolution.
"""
from __future__ import annotations

import ast
import collections
import inspect
import math
import os
import sys
import textwrap
from pathlib import Path
from types import MethodType

# Third-party code is used unmodified from source: SAM 2 (facebookresearch/sam2 @ 2b90b9f) and
# MUSt3R (naver/must3r @ 8d1fcbc, with its dust3r / croco submodules). The Space vendors them as
# sam2-src/ and must3r/ + dust3r/; the GitHub repository uses third_party/ submodules.
_ROOT = Path(__file__).resolve().parent
for _path in (_ROOT / 'third_party' / 'must3r', _ROOT / 'third_party' / 'sam2', _ROOT / 'sam2-src'):
    if _path.is_dir() and str(_path) not in sys.path:
        sys.path.insert(0, str(_path))

import numpy as np
import roma
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
from PIL import Image
from torchvision import transforms as T

from sam2.build_sam import build_sam2_video_predictor
from sam2.utils.misc import fill_holes_in_mask_scores
from sam2.modeling.sam.transformer import RoPEAttention
from sam2.modeling.position_encoding import apply_rotary_enc
from sam2.modeling.sam2_base import SAM2Base

from must3r.model import ActivationType, apply_activation, load_model, get_pointmaps_activation
from must3r.tools.image import get_HW_resolution
import must3r.tools.path_to_dust3r  # noqa: F401  registers dust3r on sys.path
from dust3r.post_process import estimate_focal_knowing_depth

MUST3R_SIZE = 512
SAM2_IMAGE_SIZE = 1024
MUST3R_LAYERS = (0, 5, 8, 12)          # encoder output + decoder blocks 5, 8, 12
MUST3R_CHANNELS = (1024, 768, 768, 768)
MEMORY_POLICY_NAME = 'previous_frame_then_recent_positive_masks'

IMAGE_TRANSFORM = T.Compose([
    T.ToTensor(),
    T.Resize((SAM2_IMAGE_SIZE, SAM2_IMAGE_SIZE), interpolation=T.InterpolationMode.NEAREST_EXACT),
    T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
])


# ----------------------------------------------------------------------------------------
# Geometry helpers (letterbox instead of crop)
# ----------------------------------------------------------------------------------------

def letterbox_geometry(width, height):
    target_h, target_w = get_HW_resolution(height, width, MUST3R_SIZE)
    padded_w = max(width, int(math.ceil(height * target_w / target_h)))
    padded_h = max(height, int(math.ceil(width * target_h / target_w)))
    return dict(width=width, height=height, padded_w=padded_w, padded_h=padded_h,
                left=(padded_w - width) // 2, top=(padded_h - height) // 2,
                target_w=target_w, target_h=target_h)


def pad_frame(rgb: np.ndarray, g: dict) -> np.ndarray:
    padded = np.zeros((g['padded_h'], g['padded_w'], 3), np.uint8)
    padded[g['top']:g['top'] + g['height'], g['left']:g['left'] + g['width']] = rgb
    return padded


def sam2_input(padded: np.ndarray) -> torch.Tensor:
    """[3, 1024, 1024] normalized SAM2 input from a padded RGB uint8 frame (nearest-exact)."""
    return IMAGE_TRANSFORM(Image.fromarray(padded))


def must3r_input(padded: np.ndarray, g: dict) -> torch.Tensor:
    """[3, target_h, target_w] in [-1, 1] for MUSt3R (nearest-exact)."""
    import cv2
    small = cv2.resize(padded, (g['target_w'], g['target_h']), interpolation=cv2.INTER_NEAREST_EXACT)
    return torch.from_numpy(small).permute(2, 0, 1).contiguous().float() / 255 * 2 - 1


def points_to_sam2(points, labels, g: dict, device):
    """Original-pixel clicks -> SAM2 1024x1024 coordinates on the padded frame."""
    coords = [[(x + g['left']) * SAM2_IMAGE_SIZE / g['padded_w'],
               (y + g['top']) * SAM2_IMAGE_SIZE / g['padded_h']] for x, y in points]
    return dict(point_coords=torch.tensor([coords], dtype=torch.float32, device=device),
                point_labels=torch.tensor([list(labels)], dtype=torch.int32, device=device))


def logits_to_original_mask(logits: torch.Tensor, g: dict) -> np.ndarray:
    """[1024, 1024] logits -> bool mask at the original frame resolution."""
    resized = F.interpolate(logits[None, None].float(), size=(g['padded_h'], g['padded_w']),
                            mode='bilinear', align_corners=False)[0, 0]
    y, x = g['top'], g['left']
    return (resized[y:y + g['height'], x:x + g['width']] > 0).cpu().numpy()


def original_mask_to_sam2(mask: np.ndarray, g: dict, device) -> torch.Tensor:
    """bool mask at original resolution -> [1, 1, 1024, 1024] bool mask input on the padded frame."""
    padded = np.zeros((g['padded_h'], g['padded_w']), np.float32)
    padded[g['top']:g['top'] + g['height'], g['left']:g['left'] + g['width']] = mask.astype(np.float32)
    t = torch.from_numpy(padded)[None, None]
    t = F.interpolate(t, size=(SAM2_IMAGE_SIZE, SAM2_IMAGE_SIZE), mode='nearest-exact')
    return (t >= 0.5).to(device)


# ----------------------------------------------------------------------------------------
# Inlined 3AM modules (identical to the training / evaluation code)
# ----------------------------------------------------------------------------------------

def positional_encoding(positions, freqs, dim=1):
    assert 0 <= dim < positions.ndim
    bands = 2 ** torch.arange(freqs, dtype=positions.dtype, device=positions.device)
    expanded = positions.unsqueeze(dim + 1) * bands.view(-1, *([1] * (positions.ndim - dim - 1)))
    expanded = expanded.reshape(*positions.shape[:dim], -1, *positions.shape[dim + 1:])
    return torch.cat([torch.sin(expanded), torch.cos(expanded), positions], dim=dim)


@torch.autocast("cuda", dtype=torch.float32)
def postprocess_must3r_output(pointmaps, pointmaps_activation=ActivationType.NORM_EXP, compute_cam=True):
    out = {}
    C = pointmaps.shape[-1]
    out['pts3d'] = apply_activation(pointmaps[..., :3], activation=pointmaps_activation)
    if C >= 6:
        out['pts3d_local'] = apply_activation(pointmaps[..., 3:6], activation=pointmaps_activation)
    if C in (4, 7):
        out['conf'] = 1.0 + pointmaps[..., -1].exp()
    if not compute_cam:
        return out
    batch_dims = out['pts3d'].shape[:-3]
    nbd = len(batch_dims)
    H, W = out['conf'].shape[-2:]
    pp = torch.tensor((W / 2, H / 2), device=out['pts3d'].device)
    focal = estimate_focal_knowing_depth(
        out['pts3d_local'].reshape(math.prod(batch_dims), H, W, 3), pp, focal_mode='weiszfeld')
    out['focal'] = focal.reshape(*batch_dims)
    R, T_ = roma.rigid_points_registration(
        out['pts3d_local'].reshape(*batch_dims, -1, 3),
        out['pts3d'].reshape(*batch_dims, -1, 3),
        weights=out['conf'].reshape(*batch_dims, -1) - 1.0, compute_scaling=False)
    c2w = torch.eye(4, device=out['pts3d'].device).view(*([1] * nbd), 4, 4).repeat(*batch_dims, 1, 1)
    c2w[..., :3, :3] = R
    c2w[..., :3, 3] = T_.view(*batch_dims, 3)
    out['c2w'] = c2w
    ys, xs = torch.meshgrid(torch.arange(H, device=out['pts3d'].device),
                            torch.arange(W, device=out['pts3d'].device), indexing='ij')
    f = out['focal'].reshape(*batch_dims, 1, 1)
    x = (xs - pp[0]) / f
    y = (ys - pp[1]) / f
    d_cam = F.normalize(torch.stack([x, y, torch.ones_like(x)], dim=-1), dim=-1)
    d_world = torch.einsum('...ij,...hwj->...hwi', R, d_cam)
    o_world = c2w[..., :3, 3].view(*batch_dims, 1, 1, 3).expand(*batch_dims, H, W, 3)
    out['ray_origin'] = o_world
    out['ray_dir'] = d_world
    out['ray_plucker'] = torch.cat([torch.cross(o_world, d_world, dim=-1), d_world], dim=-1)
    return out


class LayerNorm2d(nn.Module):
    def __init__(self, num_channels: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(num_channels))
        self.bias = nn.Parameter(torch.zeros(num_channels))
        self.eps = eps

    def forward(self, x):
        u = x.mean(1, keepdim=True)
        s = (x - u).pow(2).mean(1, keepdim=True)
        x = (x - u) / torch.sqrt(s + self.eps)
        return self.weight[:, None, None] * x + self.bias[:, None, None]


class RoutingFusion(nn.Module):
    """Additive FiLM-gated fusion of MUSt3R decoder features + explicit 3D geometry into feat_2d."""
    def __init__(self, d_2d=256, d_3d=(1024, 768, 768, 768), d_route=256, d_hidden=384, explicit_3d_freqs=6):
        super().__init__()
        self.L = len(d_3d)
        self.d_2d = d_2d
        self.d_route = d_route
        self.explicit_3d_freqs = explicit_3d_freqs
        explicit_3d_in_ch = 3 * (2 * explicit_3d_freqs + 1) + 6
        self.explicit_3d_embedding = nn.Conv2d(explicit_3d_in_ch, d_2d, kernel_size=16, stride=16)
        self.proj = nn.ModuleList([
            nn.Sequential(nn.Conv2d(c, d_hidden, kernel_size=1, bias=True), LayerNorm2d(d_hidden), nn.GELU(),
                          nn.Conv2d(d_hidden, d_2d, kernel_size=3, padding=1, bias=True))
            for c in d_3d])
        self.mod_mlp = nn.ModuleList([
            nn.Sequential(nn.Linear(d_route, 2 * d_route), nn.GELU(), nn.Linear(2 * d_route, 2 * d_2d + 1))
            for _ in range(self.L + 1)])
        self.base_mod_mlp = nn.Sequential(nn.Linear(d_route, 2 * d_route), nn.GELU(), nn.Linear(2 * d_route, 2 * d_2d))
        nn.init.zeros_(self.base_mod_mlp[-1].weight)
        nn.init.zeros_(self.base_mod_mlp[-1].bias)

    def _film_add(self, feat, delta, mod, C, per_layer_alpha):
        gamma = 1.0 + mod[:, :C, None, None]
        beta = mod[:, C:2 * C, None, None]
        alpha = mod[:, 2 * C:, None, None]
        feat = feat + alpha * (delta * gamma + beta)
        per_layer_alpha.append(alpha.view(feat.shape[0]))
        return feat

    def forward(self, feat_2d, feat_3d_list, r, explicit_3d):
        B, C, H, W = feat_2d.shape
        assert r.shape[1] == self.L + 2
        base_mod = self.base_mod_mlp(r[:, self.L + 1, :]).to(dtype=feat_2d.dtype)
        feat = feat_2d * (1.0 + base_mod[:, :C, None, None]) + base_mod[:, C:2 * C, None, None]
        per_layer_alpha = []
        for l in range(self.L):
            delta = self.proj[l](feat_3d_list[l])
            delta = F.interpolate(delta, size=(H, W), mode='bilinear', align_corners=False)
            mod = self.mod_mlp[l](r[:, l, :]).to(dtype=feat.dtype)
            feat = self._film_add(feat, delta, mod, C, per_layer_alpha)
        pts3d = explicit_3d[:, :3]
        plucker = explicit_3d[:, 3:]
        pe_in = torch.cat([positional_encoding(pts3d, self.explicit_3d_freqs, dim=1), plucker], dim=1)
        delta = self.explicit_3d_embedding(pe_in.to(self.explicit_3d_embedding.weight.dtype))
        delta = F.interpolate(delta, size=(H, W), mode='bilinear', align_corners=False).to(feat.dtype)
        mod = self.mod_mlp[self.L](r[:, self.L, :]).to(dtype=feat.dtype)
        feat = self._film_add(feat, delta, mod, C, per_layer_alpha)
        alpha_stack = torch.stack(per_layer_alpha, dim=1)
        return feat, alpha_stack, None, base_mod.detach()


def _cross_attn_with_routing(ca, full_q, mem_k_in, mem_v_in, HW, num_k_exclude_rope=0):
    q = ca.q_proj(full_q)
    k = ca.k_proj(mem_k_in)
    v = ca.v_proj(mem_v_in)
    q = ca._separate_heads(q, ca.num_heads)
    k = ca._separate_heads(k, ca.num_heads)
    v = ca._separate_heads(v, ca.num_heads)
    if isinstance(ca, RoPEAttention):
        freqs_cis = ca.freqs_cis.to(q.device)
        assert freqs_cis.shape[0] == HW
        num_k_rope = k.size(-2) - num_k_exclude_rope
        q_spatial = q[:, :, :HW]
        q_routing = q[:, :, HW:]
        q_spatial, k_rope = apply_rotary_enc(q_spatial, k[:, :, :num_k_rope], freqs_cis=freqs_cis,
                                             repeat_freqs_k=ca.rope_k_repeat)
        q = torch.cat([q_spatial, q_routing], dim=2)
        k = torch.cat([k_rope, k[:, :, num_k_rope:]], dim=2)
    dropout_p = ca.dropout_p if ca.training else 0.0
    out = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p)
    out = ca._recombine_heads(out)
    return ca.out_proj(out)


class RoutingMemoryAttention(nn.Module):
    """Replaces SAM2's MemoryAttention: routing tokens ride inside every layer and see the memory bank."""
    def __init__(self, base, fusion: RoutingFusion, d_route=256, routing_num_heads=8, num_routing_tokens_per_layer=2):
        super().__init__()
        self.base = base
        self.fusion = fusion
        self.d_route = d_route
        self.L = fusion.L
        self.G = self.L + 2
        self.K = num_routing_tokens_per_layer
        self.N = self.K * self.G
        d_model = base.d_model
        kv_in_dim = base.layers[0].cross_attn_image.kv_in_dim
        self.routing_token = nn.Parameter(torch.zeros(1, self.N, d_model))
        self.r_in_proj = nn.Linear(d_route, d_model)
        self.r_out_proj = nn.Linear(self.K * d_model, d_route)
        self.r0 = nn.Parameter(torch.zeros(1, self.G, d_route))
        self._expected_obj_ptr_tokens = max(1, d_model // kv_in_dim)
        self.obj_ptr_to_r = nn.Linear(self._expected_obj_ptr_tokens * kv_in_dim, self.G * d_route)
        nn.init.zeros_(self.obj_ptr_to_r.weight)
        nn.init.zeros_(self.obj_ptr_to_r.bias)
        nn.init.zeros_(self.r_out_proj.weight)
        nn.init.zeros_(self.r_out_proj.bias)
        self._feat_2d_seq = None
        self._feat_3d_seq = None
        self._explicit_3d_seq = None
        self._frame_idx = 1
        self._r = None
        self._feats_for_memory_out = None
        self.history = {}

    @property
    def num_layers(self):
        return self.base.num_layers

    @property
    def d_model(self):
        return self.base.d_model

    def reset_state(self, feat_2d_seq, feat_3d_seq, explicit_3d_seq, B, device, dtype, feats_for_memory_out=None):
        self._feat_2d_seq = feat_2d_seq
        self._feat_3d_seq = feat_3d_seq
        self._explicit_3d_seq = explicit_3d_seq
        self._frame_idx = 1
        self._r = self.r0.expand(B, -1, -1).to(device=device, dtype=dtype)
        self._feats_for_memory_out = feats_for_memory_out
        self.history = {'r': [self._r.detach()], 'alpha': [], 'base_mod': []}

    def clear_state(self):
        self._feat_2d_seq = None
        self._feat_3d_seq = None
        self._explicit_3d_seq = None
        self._frame_idx = 1
        self._r = None
        self._feats_for_memory_out = None
        self.history = {}

    def forward(self, curr, memory, curr_pos=None, memory_pos=None, num_obj_ptr_tokens=0):
        curr_t = curr[0] if isinstance(curr, list) else curr
        curr_pos_t = curr_pos[0] if isinstance(curr_pos, list) else curr_pos
        base = self.base
        B_size = curr_t.shape[1]
        if self._frame_idx == 1 and num_obj_ptr_tokens > 0:
            assert num_obj_ptr_tokens == self._expected_obj_ptr_tokens
            obj_ptr_seq = memory[-num_obj_ptr_tokens:]
            obj_ptr_recon = obj_ptr_seq.transpose(0, 1).flatten(1, 2)
            r_delta = self.obj_ptr_to_r(obj_ptr_recon).view(B_size, self.G, self.d_route)
            self._r = self._r + r_delta.to(self._r.dtype)
        t = min(self._frame_idx, len(self._feat_2d_seq) - 1)
        feat_2d = self._feat_2d_seq[t].to(device=curr_t.device, dtype=curr_t.dtype)
        feat_3d = [f.to(device=curr_t.device, dtype=curr_t.dtype) for f in self._feat_3d_seq[t]]
        explicit_3d = self._explicit_3d_seq[t].to(device=curr_t.device, dtype=curr_t.dtype)
        fused, alpha, _, base_mod = self.fusion(feat_2d, feat_3d, self._r, explicit_3d)
        fused_flat = rearrange(fused, 'b c h w -> (h w) b c')
        HW, B, C = fused_flat.shape
        N, K, L, G = self.N, self.K, self.L, self.G
        if self._feats_for_memory_out is not None and t < len(self._feats_for_memory_out):
            self._feats_for_memory_out[t] = [fused_flat]
        r_inject = self.r_in_proj(self._r).to(curr_t.dtype)
        r_inject = r_inject.unsqueeze(2).expand(B, G, K, C).reshape(B, N, C)
        rt = self.routing_token.expand(1, N, C).repeat(B, 1, 1)
        rt = rt + r_inject
        spatial = fused_flat
        if base.pos_enc_at_input and curr_pos_t is not None:
            spatial = spatial + 0.1 * curr_pos_t
        assert base.batch_first
        spatial = spatial.transpose(0, 1)
        spatial_pos = curr_pos_t.transpose(0, 1)
        mem_b = memory.transpose(0, 1)
        mem_pos_b = memory_pos.transpose(0, 1)
        rt_pos = torch.zeros_like(rt)
        for layer in base.layers:
            tgt = torch.cat([spatial, rt], dim=1)
            qpos = torch.cat([spatial_pos, rt_pos], dim=1)
            tgt2 = layer.norm1(tgt)
            full_q = tgt2 + qpos if layer.pos_enc_at_attn else tgt2
            tgt2 = _cross_attn_with_routing(layer.self_attn, full_q, full_q, tgt2, HW=HW, num_k_exclude_rope=N)
            tgt = tgt + layer.dropout1(tgt2)
            tgt2 = layer.norm2(tgt)
            full_q = tgt2 + qpos if layer.pos_enc_at_cross_attn_queries else tgt2
            mem_k_in = mem_b + mem_pos_b if layer.pos_enc_at_cross_attn_keys else mem_b
            ca_out = _cross_attn_with_routing(layer.cross_attn_image, full_q, mem_k_in, mem_b, HW=HW,
                                              num_k_exclude_rope=num_obj_ptr_tokens)
            tgt = tgt + layer.dropout2(ca_out)
            tgt2 = layer.norm3(tgt)
            tgt2 = layer.linear2(layer.dropout(layer.activation(layer.linear1(tgt2))))
            tgt = tgt + layer.dropout3(tgt2)
            spatial, rt = tgt[:, :HW], tgt[:, HW:]
        spatial = base.norm(spatial)
        rt = base.norm(rt)
        spatial = spatial.transpose(0, 1)
        rt_per_layer = rt.reshape(B, G, K * C)
        new_r = self.r_out_proj(rt_per_layer).to(self._r.dtype)
        self._r = new_r
        self._frame_idx += 1
        self.history['r'].append(new_r.detach())
        self.history['alpha'].append(alpha)
        self.history['base_mod'].append(base_mod)
        return spatial


# ----------------------------------------------------------------------------------------
# SAM2 memory policy: previous frame, then the most recent nonempty masks
# ----------------------------------------------------------------------------------------

def _previous_visible_indices(outputs, frame_idx):
    previous = frame_idx - 1
    keys = [previous] if previous in outputs else []
    for k in sorted(outputs, reverse=True):
        if k >= previous:
            continue
        out = outputs[k]
        if '_memory_visible' not in out:
            out['_memory_visible'] = bool((out['pred_masks'] > 0).any())
        if out['_memory_visible']:
            keys.append(k)
    return keys


def _build_previous_visible_method(base_cls):
    """AST-patch SAM2Base._prepare_memory_conditioned_features so that only the memory *selection*
    changes (previous frame + recent positive masks, real pointer distances); encodings stay official."""
    original = base_cls._prepare_memory_conditioned_features
    tree = ast.parse(textwrap.dedent(inspect.getsource(original)))
    fn = tree.body[0]
    fn.name = '_prepare_previous_visible'
    fn.body.insert(1, ast.parse("assert not track_in_reverse, 'Harness uses forward-indexed passes'").body[0])
    fn.body.insert(2, ast.parse("visible_keys = _previous_visible_indices(output_dict['non_cond_frame_outputs'], frame_idx)").body[0])
    counts = {'spatial': 0, 'pointers': 0}

    class Patch(ast.NodeTransformer):
        def visit_Assign(self, node):
            self.generic_visit(node)
            if (len(node.targets) == 1 and isinstance(node.targets[0], ast.Name) and node.targets[0].id == 'out'
                    and isinstance(node.value, ast.Call) and isinstance(node.value.func, ast.Attribute)
                    and node.value.func.attr == 'get' and isinstance(node.value.func.value, ast.Subscript)
                    and node.value.args and isinstance(node.value.args[0], ast.Name)
                    and node.value.args[0].id == 'prev_frame_idx'):
                counts['spatial'] += 1
                override = ast.parse('prev_frame_idx = visible_keys[t_rel-1] if t_rel <= len(visible_keys) else -1').body[0]
                return [override, node]
            return node

        def visit_For(self, node):
            self.generic_visit(node)
            if isinstance(node.target, ast.Name) and node.target.id == 't_diff':
                counts['pointers'] += 1
                replacement = ast.parse('for t_diff, t in [(frame_idx-k,k) for k in visible_keys[:max_obj_ptrs_in_encoder-1]]:\n pass').body[0]
                assert isinstance(node.body[0], ast.Assign) and node.body[0].targets[0].id == 't'
                replacement.body = node.body[1:]
                return replacement
            return node

    tree = Patch().visit(tree)
    ast.fix_missing_locations(tree)
    assert counts == {'spatial': 1, 'pointers': 1}, counts
    namespace = dict(original.__globals__)
    namespace['_previous_visible_indices'] = _previous_visible_indices
    exec(compile(tree, '<previous-visible>', 'exec'), namespace)
    return namespace['_prepare_previous_visible']


def install_memory_policy(predictor):
    method = _build_previous_visible_method(SAM2Base)
    predictor._prepare_memory_conditioned_features = MethodType(method, predictor)
    predictor.memory_temporal_stride_for_eval = 1


class PositiveMemory:
    """Keep useful (nonempty) outputs across arbitrarily long gaps of empty masks."""
    def __init__(self, capacity):
        self.capacity = capacity
        self.outputs = {}

    def restore(self, noncond):
        noncond.update(self.outputs)

    def remember(self, index, output):
        output['_memory_visible'] = bool((output['pred_masks'] > 0).any())
        if output['_memory_visible']:
            self.outputs[index] = output
        while len(self.outputs) > self.capacity:
            del self.outputs[min(self.outputs)]


# ----------------------------------------------------------------------------------------
# Model construction
# ----------------------------------------------------------------------------------------

import contextlib


@contextlib.contextmanager
def cpu_construction():
    """Build models without touching CUDA. SAM2 moves RoPE tables / position caches to 'cuda' at
    construction when torch.cuda.is_available(); on ZeroGPU that happens at import time and the
    resulting placeholder tensors cannot be deep-copied next to CPU tensors. Forward passes move
    those tables to the input device anyway."""
    original = torch.cuda.is_available
    torch.cuda.is_available = lambda: False
    try:
        yield
    finally:
        torch.cuda.is_available = original


def build_3am_predictor(checkpoint_path, device='cpu'):
    with cpu_construction():
        predictor = build_sam2_video_predictor('configs/sam2.1/sam2.1_hiera_l.yaml', device='cpu',
                                               apply_postprocessing=True).eval()
    d_route = predictor.memory_attention.d_model
    fusion = RoutingFusion(d_2d=256, d_3d=(1024, 768, 768, 768), d_route=d_route)
    predictor.memory_attention = RoutingMemoryAttention(base=predictor.memory_attention, fusion=fusion,
                                                        d_route=d_route, num_routing_tokens_per_layer=2).to(device).eval()
    weights = torch.load(checkpoint_path, map_location='cpu', mmap=True, weights_only=False)
    predictor.load_state_dict(weights.get('model', weights), strict=True)
    del weights
    install_memory_policy(predictor)
    return predictor.to(device)


def build_original_predictor(checkpoint_path, device='cpu'):
    with cpu_construction():
        predictor = build_sam2_video_predictor('configs/sam2.1/sam2.1_hiera_l.yaml', checkpoint_path,
                                               device='cpu', apply_postprocessing=True).eval()
    return predictor.to(device)


def build_must3r(weights_path, device='cpu'):
    from must3r.model.blocks.attention import has_xformers, toggle_memory_efficient_attention
    # xformers is off unless MUST3R_XFORMERS=1: on the Space the wheel has no CUDA ops for the runtime
    # GPU, and PyTorch SDPA gives the same features (verified against the reference implementation).
    toggle_memory_efficient_attention(enabled=bool(has_xformers and os.environ.get('MUST3R_XFORMERS', '0') == '1'))
    with cpu_construction():
        model = load_model(weights_path, device='cpu', img_size=MUST3R_SIZE, memory_mode='kv', verbose=False)
    model = [m.eval().to(device) for m in model]
    assert model[0].patch_size == 16
    assert get_pointmaps_activation(model[1]) == ActivationType.NORM_EXP
    return model


# ----------------------------------------------------------------------------------------
# Feature extraction
# ----------------------------------------------------------------------------------------

@torch.no_grad()
@torch.autocast(device_type='cuda', dtype=torch.bfloat16)
def get_image_feature(predictor, images):
    backbone_out = predictor.forward_image(images)
    backbone_out = {'backbone_fpn': backbone_out['backbone_fpn'].copy(),
                    'vision_pos_enc': backbone_out['vision_pos_enc'].copy()}
    backbone_out, vision_feats, vision_pos_embeds, feat_sizes = predictor._prepare_backbone_features(backbone_out)
    return vision_feats, vision_pos_embeds, feat_sizes


def must3r_frames(padded_frames, g) -> np.ndarray:
    """CPU: [N, target_h, target_w, 3] uint8 MUSt3R inputs (nearest-exact, as must3r_input)."""
    import cv2
    return np.stack([cv2.resize(p, (g['target_w'], g['target_h']), interpolation=cv2.INTER_NEAREST_EXACT)
                     for p in padded_frames])


def sam2_frames(padded_frames) -> np.ndarray:
    """CPU: [N, 1024, 1024, 3] uint8 SAM2 inputs. Nearest-exact selects pixels, so resizing the uint8
    frame and normalizing later on the GPU gives exactly sam2_input()."""
    out = np.empty((len(padded_frames), SAM2_IMAGE_SIZE, SAM2_IMAGE_SIZE, 3), np.uint8)
    for i, p in enumerate(padded_frames):
        t = torch.from_numpy(np.ascontiguousarray(p)).permute(2, 0, 1)[None]
        out[i] = F.interpolate(t, size=(SAM2_IMAGE_SIZE, SAM2_IMAGE_SIZE), mode='nearest-exact')[0].permute(1, 2, 0).numpy()
    return out


_SAM2_MEAN = torch.tensor((0.485, 0.456, 0.406)).view(1, 3, 1, 1)
_SAM2_STD = torch.tensor((0.229, 0.224, 0.225)).view(1, 3, 1, 1)


def u8_to_unit(x):
    """uint8 tensor -> float in [0, 1], bit-identical to the CPU path. With a Python-scalar divisor
    CUDA multiplies by the reciprocal, which differs by one ulp and is amplified by the memories."""
    return x.float() / torch.full((), 255.0, device=x.device)


def sam2_batch(frames_u8, device):
    """[B, 1024, 1024, 3] uint8 (numpy) -> normalized [B, 3, 1024, 1024] float on `device`."""
    x = u8_to_unit(torch.from_numpy(np.ascontiguousarray(frames_u8)).to(device).permute(0, 3, 1, 2))
    return (x - _SAM2_MEAN.to(device)) / _SAM2_STD.to(device)


class _LayerTap:
    """Capture per-view MUSt3R features at MUST3R_LAYERS with forward hooks, so stock MUSt3R needs no
    changes: layer 0 is the encoder token input of feat_embed_enc_to_dec, layer k is the output of
    decoder block k (before the head's final norm). The decoder visits views in order, so the i-th
    call of each hook within one decoder call belongs to view i."""

    def __init__(self, decoder, layers=MUST3R_LAYERS):
        self.layers = tuple(layers)
        self.calls = {k: [] for k in self.layers}
        self.handles = []
        for k in self.layers:
            if k == 0:
                hook = (lambda k: lambda module, args, out: self.calls[k].append(args[0]))(k)
                self.handles.append(decoder.feat_embed_enc_to_dec.register_forward_hook(hook))
            else:
                hook = (lambda k: lambda module, args, out: self.calls[k].append(out))(k)
                self.handles.append(decoder.blocks_dec[k - 1].register_forward_hook(hook))

    def pop_views(self):
        """[view][layer] tensors of shape [Ni, C] recorded since the last call."""
        counts = {len(v) for v in self.calls.values()}
        assert len(counts) == 1, counts
        views = [[self.calls[k][v] for k in self.layers] for v in range(counts.pop())]
        assert all(t.shape[0] == 1 for view in views for t in view), 'one view per decoder stack (max_bs=1)'
        for v in self.calls.values():
            v.clear()
        return [[t[0] for t in view] for view in views]

    def remove(self):
        for h in self.handles:
            h.remove()


def auto_memory_stride(n, memory_frames=300):
    """MUSt3R memory write stride for an n-frame video: every frame up to `memory_frames` frames (the
    evaluation setting), then the smallest stride that keeps the memory at about `memory_frames` frames,
    so memory and per-frame cost stay bounded on videos of any length."""
    return max(1, math.ceil(n / memory_frames))


def must3r_storage(n, g, root=None):
    """Arrays for the per-frame MUSt3R outputs (float16): feats 4 x [N, C, h, w] and explicit
    [N, 9, target_h, target_w], about 8.6 MB per frame. In host memory, or with `root` as .npy files
    (filled frame by frame, then memory-mapped by later runs) so host memory does not grow with the
    video length either."""
    h, w = g['target_h'] // 16, g['target_w'] // 16
    shapes = {f'feat{k}': (n, c, h, w) for k, c in zip(MUST3R_LAYERS, MUST3R_CHANNELS)}
    shapes['explicit'] = (n, 9, g['target_h'], g['target_w'])
    if root is None:
        arrays = {k: np.empty(s, np.float16) for k, s in shapes.items()}
    else:
        arrays = {k: np.lib.format.open_memmap(os.path.join(root, f'{k}.npy'), mode='w+', dtype=np.float16, shape=s)
                  for k, s in shapes.items()}
    return [arrays[f'feat{k}'] for k in MUST3R_LAYERS], arrays['explicit']


def load_must3r_storage(root):
    """Memory-map the arrays that run_must3r_iter filled in `root` (see must3r_storage)."""
    feats = [np.load(os.path.join(root, f'feat{k}.npy'), mmap_mode='r') for k in MUST3R_LAYERS]
    return feats, np.load(os.path.join(root, 'explicit.npy'), mmap_mode='r')


@torch.no_grad()
def run_must3r(must3r, frames_u8, g, device='cuda', memory_stride=1, max_memory_frames=0, progress=None, out=None):
    """Stream MUSt3R-512 over the frames with an explicit KV memory policy (no keyframe search,
    no KD-tree, no retrieval).

    Frames 0 and 1 initialize the memory together and always stay in it. Every later frame reads
    the current memory; its own KV is written to the memory only if `index % memory_stride == 0`.
    With `max_memory_frames > 0` the memory holds at most that many frames: the two anchors plus the
    most recent written frames (oldest evicted first), which bounds GPU memory and per-frame
    attention cost on long videos. memory_stride=1, max_memory_frames=0 keeps every frame's KV
    (the evaluation setting).

    frames_u8: [N, target_h, target_w, 3] uint8 from must3r_frames(). out: optional arrays from
    must3r_storage() to fill (default: allocated in host memory).
    Returns must3r_feats: 4 x [N, C, h, w] (float16 numpy), explicit: [N, 9, target_h, target_w] (float16 numpy), stats.
    """
    result = None
    for event in run_must3r_iter(must3r, frames_u8, g, device, memory_stride, max_memory_frames, out):
        if event[0] == 'progress':
            if progress is not None:
                progress(event[1], event[2])
        else:
            result = event[1:]
    return result


@torch.no_grad()
def run_must3r_iter(must3r, frames_u8, g, device='cuda', memory_stride=1, max_memory_frames=0, out=None):
    """Generator form of run_must3r: yields ('progress', done, total) after every frame, then
    ('result', must3r_feats, explicit, stats). Every frame's outputs are copied to `out` as soon as
    they are computed, so GPU memory holds only the models and the MUSt3R memory. Autocast is entered
    per frame, so the caller never runs under it while the generator is suspended."""
    from collections import deque
    from must3r.engine.inference import (encoder_multi_ar, inference_multi_ar_batch, stack_views,
                                         unstack_pointmaps, get_Nmem, _remove_from_mem)
    encoder, decoder = must3r
    n = int(frames_u8.shape[0])
    stride = max(1, int(memory_stride))
    cap = int(max_memory_frames)
    if 0 < cap < 3:
        raise ValueError('max_memory_frames must be 0 (keep all) or at least 3')
    h, w = g['target_h'] // 16, g['target_w'] // 16
    shape = torch.tensor([g['target_h'], g['target_w']], dtype=torch.int32, device=device)
    pointmaps_activation = get_pointmaps_activation(decoder, verbose=False)
    if hasattr(decoder, 'all_feats'):
        del decoder.all_feats   # older recording hook: would copy all 13 layers of every frame to the host
    feats, explicit = out if out is not None else must3r_storage(n, g)
    mem, written, peak = None, deque(), 0
    batches = [list(range(min(2, n)))] + [[i] for i in range(2, n)]
    tap = _LayerTap(decoder)
    try:
        for ids in batches:
            with torch.autocast('cuda', dtype=torch.bfloat16):
                imgs = u8_to_unit(torch.from_numpy(np.ascontiguousarray(frames_u8[ids[0]:ids[-1] + 1])).to(device)
                                  .permute(0, 3, 1, 2)) * 2 - 1
                imgs = list(imgs)
                true_shape = torch.stack([shape] * len(ids))
                x, pos = encoder_multi_ar(encoder, imgs, true_shape, max_bs=1, device=device)
                shape_stacks, index_stacks, x_stacks, pos_stacks, img_stacks = stack_views(
                    true_shape, [x, pos, imgs], max_bs=1)
                n_before = get_Nmem(mem)
                new_mem, pointmaps = inference_multi_ar_batch(
                    encoder, decoder, img_stacks, shape_stacks, mem, encoder_precomputed_features=(x_stacks, pos_stacks),
                    post_process_function=lambda x: {'pm': x}, device=device)
                pointmaps = unstack_pointmaps(index_stacks, pointmaps)
                layers = [None] * len(ids)
                for (view,), view_layers in zip(index_stacks, tap.pop_views()):
                    layers[view] = view_layers
                new_mem = list(new_mem)
                if mem is None:
                    mem = new_mem                                   # anchors: frames 0 and 1
                elif ids[0] % stride == 0:
                    label = int(new_mem[1][0, n_before:].unique().item())
                    mem = new_mem
                    written.append(label)
                    while cap and len(written) + min(2, n) > cap:
                        mem[0], mem[1] = _remove_from_mem(mem[0], mem[1], written.popleft())
                # else: the frame read the memory but is not written to it (keep the previous memory)
                peak = max(peak, get_Nmem(mem))
                for j, i in enumerate(ids):
                    for slot, t in enumerate(layers[j]):
                        feats[slot][i] = rearrange(t, '(h w) c -> c h w', h=h, w=w).to(torch.float16).cpu().numpy()
                    pm = postprocess_must3r_output(pointmaps[j]['pm'][None], pointmaps_activation=pointmaps_activation,
                                                   compute_cam=True)
                    explicit[i] = (torch.cat((pm['pts3d'], pm['ray_plucker']), dim=-1).squeeze(0).permute(2, 0, 1)
                                   .to(torch.float16).cpu().numpy())
            yield ('progress', ids[-1] + 1, n)
        tokens_per_frame = h * w
        stats = dict(frames=n, memory_stride=stride, max_memory_frames=cap,
                     memory_frames_final=get_Nmem(mem) // tokens_per_frame, memory_frames_peak=peak // tokens_per_frame)
    finally:
        tap.remove()
        del mem
        torch.cuda.empty_cache()
    yield ('result', feats, explicit, stats)


# ----------------------------------------------------------------------------------------
# Tracking (evaluation tracker + positive-memory bank)
# ----------------------------------------------------------------------------------------

class Tracker3AM(nn.Module):
    def __init__(self, predictor, device='cuda'):
        super().__init__()
        self.predictor = predictor
        self.device = device

    @torch.no_grad()
    @torch.autocast(device_type='cuda', dtype=torch.bfloat16)
    def initialize(self, vision_feats_all, vision_pos_all, feat_sizes, must3r_feats, explicit_3d,
                   processing_order, mask_inputs=None, point_inputs=None):
        device, bf16 = self.device, torch.bfloat16
        order = list(processing_order)
        N = len(order)
        # per-frame inputs of the routing fusion, fetched when the tracker reaches the frame
        feat_2d_seq = _LazySeq(lambda t: rearrange(vision_feats_all[order[t]][-1], '(h w) b c -> b c h w', h=64, w=64).contiguous(), N)
        feat_3d_seq = _LazySeq(lambda t: [_frame(f, order[t], device) for f in must3r_feats], N)
        explicit_3d_seq = _LazySeq(lambda t: _frame(explicit_3d, order[t], device), N)
        r0 = self.predictor.memory_attention.r0.expand(1, -1, -1).to(device=device, dtype=bf16)
        fused_0, _, _, _ = self.predictor.memory_attention.fusion(
            feat_2d_seq[0].to(device=device, dtype=bf16), [f.to(device=device, dtype=bf16) for f in feat_3d_seq[0]],
            r0, explicit_3d_seq[0].to(device=device, dtype=bf16))
        feats_for_memory = [None] * N
        feats_for_memory[0] = [rearrange(fused_0, 'b c h w -> (h w) b c')]
        self.predictor.memory_attention.reset_state(feat_2d_seq=feat_2d_seq, feat_3d_seq=feat_3d_seq,
                                                    explicit_3d_seq=explicit_3d_seq, B=1, device=device, dtype=bf16,
                                                    feats_for_memory_out=feats_for_memory)
        self.processing_order, self.N = order, N
        self.mask_inputs = mask_inputs.to(device) if mask_inputs is not None else None
        self.point_inputs = point_inputs
        self.fused_0 = fused_0
        self.feats_for_memory = feats_for_memory
        self.vision_feats_all, self.vision_pos_all, self.feat_sizes = vision_feats_all, vision_pos_all, feat_sizes
        self.output_dict = {'cond_frame_outputs': {}, 'non_cond_frame_outputs': {}}
        self.current_idx = 0
        self.positive_memory = PositiveMemory(max(32, self.predictor.num_maskmem, self.predictor.max_obj_ptrs_in_encoder))

    @torch.no_grad()
    @torch.autocast(device_type='cuda', dtype=torch.bfloat16)
    def step(self):
        device = self.device
        t = self.current_idx
        orig_idx = self.processing_order[t]
        self.positive_memory.restore(self.output_dict['non_cond_frame_outputs'])
        vf = [v.to(device) for v in self.vision_feats_all[orig_idx]]
        vpe = [v.to(device) for v in self.vision_pos_all[orig_idx]]
        if t == 0:
            vf[-1] = rearrange(self.fused_0, 'b c h w -> (h w) b c').contiguous()
        current_out = self.predictor.track_step(
            frame_idx=t, is_init_cond_frame=(t == 0), current_vision_feats=vf, current_vision_pos_embeds=vpe,
            feat_sizes=self.feat_sizes,
            point_inputs=self.point_inputs if t == 0 else None,
            mask_inputs=self.mask_inputs if t == 0 else None,
            output_dict=self.output_dict, num_frames=self.N, track_in_reverse=False,
            run_mem_encoder=False, prev_sam_mask_logits=None)
        memory_high_res = current_out['pred_masks_high_res']
        if self.predictor.fill_hole_area > 0:
            current_out['pred_masks'] = fill_holes_in_mask_scores(current_out['pred_masks'], self.predictor.fill_hole_area)
        current_out['pred_masks_high_res'] = F.interpolate(current_out['pred_masks'],
                                                           size=(self.predictor.image_size, self.predictor.image_size),
                                                           mode='bilinear', align_corners=False)
        if t == 0:
            memory_high_res = current_out['pred_masks_high_res']
        mem_vf = list(vf)
        mem_vf[-1] = self.feats_for_memory[t][0]
        maskmem_features, maskmem_pos_enc = self.predictor._encode_new_memory(
            current_vision_feats=mem_vf, feat_sizes=self.feat_sizes, pred_masks_high_res=memory_high_res,
            object_score_logits=current_out['object_score_logits'], is_mask_from_pts=(t == 0))
        current_out['maskmem_features'] = maskmem_features.to(torch.bfloat16)
        current_out['maskmem_pos_enc'] = maskmem_pos_enc
        logits = current_out['pred_masks_high_res'].squeeze().float()   # stays on the device
        bucket = 'cond_frame_outputs' if t == 0 else 'non_cond_frame_outputs'
        self.output_dict[bucket][t] = current_out
        self.feats_for_memory[t] = None
        noncond = self.output_dict['non_cond_frame_outputs']
        if t > 0:
            self.positive_memory.remember(t, noncond[t])
        for index in list(noncond):
            noncond[index].pop('pred_masks_high_res', None)
            if index < t - 32:
                del noncond[index]
        for values in self.predictor.memory_attention.history.values():
            values.clear()
        self.current_idx += 1
        return logits


class _LazySeq:
    """Read-only sequence whose items are computed when they are accessed."""

    def __init__(self, item, n):
        self.item, self.n = item, n

    def __len__(self):
        return self.n

    def __getitem__(self, t):
        return self.item(t)


def _frame(x, i, device):
    """Frame i of a per-frame array (numpy, possibly memory-mapped, or torch) as a [1, ...] tensor on `device`."""
    t = torch.from_numpy(np.array(x[i:i + 1])) if isinstance(x, np.ndarray) else x[i:i + 1]   # copy: memmaps are read-only
    return t.to(device)


class Sam2Features:
    """SAM2 image-encoder features per frame, computed on demand. Frames are encoded in the same aligned
    batches as encoding every frame up front (so the features are identical), and only the last few
    batches are kept, so GPU memory does not grow with the video length. frames_u8 from sam2_frames()."""

    def __init__(self, predictor, frames_u8, device='cuda', batch=4, keep=3):
        self.predictor, self.frames_u8, self.device, self.batch, self.keep = predictor, frames_u8, device, batch, keep
        self.batches = collections.OrderedDict()
        self.pos = self.feat_sizes = None   # positional embeddings (the same for every frame) and feature sizes

    def __len__(self):
        return len(self.frames_u8)

    def __getitem__(self, i):
        b, k = divmod(int(i), self.batch)
        if b in self.batches:
            self.batches.move_to_end(b)
        else:
            with torch.autocast('cuda', enabled=False):   # the image encoder runs in full precision
                images = sam2_batch(self.frames_u8[b * self.batch:(b + 1) * self.batch], self.device)
                vf, vpe, fs = get_image_feature(self.predictor, images)
            self.batches[b] = [[v[:, j:j + 1] for v in vf] for j in range(images.shape[0])]
            if self.pos is None:
                self.pos, self.feat_sizes = [v[:, 0:1] for v in vpe], fs
            while len(self.batches) > self.keep:
                self.batches.popitem(last=False)
        return self.batches[b][k]


@torch.no_grad()
def track(predictor, frames_u8, g, must3r_feats, explicit, start_idx, first_mask=None, points=None, labels=None,
          device='cuda', progress=None, encoder_batch=4):
    """Bidirectional tracking from start_idx. frames_u8: SAM2 inputs from sam2_frames().
    Prompt: first_mask (bool, original resolution) or clicks.
    Returns {frame_idx: bool mask at original resolution}."""
    masks = {}
    for i, mask, done, total in track_iter(predictor, frames_u8, g, must3r_feats, explicit, start_idx, first_mask,
                                           points, labels, device, encoder_batch):
        masks[i] = mask
        if progress is not None and (done % 10 == 0 or done == total):
            progress(done, total)
    return masks


@torch.no_grad()
def track_iter(predictor, frames_u8, g, must3r_feats, explicit, start_idx, first_mask=None, points=None, labels=None,
               device='cuda', encoder_batch=4):
    """Generator form of track: yields (frame_idx, mask, steps_done, steps_total) as each frame is
    tracked (forward from start_idx, then backward). The prompt frame is reported once. SAM2 features
    are computed and MUSt3R outputs are moved to the GPU as tracking reaches each frame, so GPU memory
    does not grow with the video length."""
    n = len(frames_u8)
    features = Sam2Features(predictor, frames_u8, device, batch=encoder_batch)
    features[start_idx]   # the first batch also provides the feature sizes and positional embeddings
    vision_pos_all = _LazySeq(lambda i: features.pos, n)
    mask_inputs = original_mask_to_sam2(first_mask, g, device) if first_mask is not None else None
    point_inputs = points_to_sam2(points, labels, g, device) if points else None
    tracker = Tracker3AM(predictor, device)
    seen, done = set(), 0
    orders = [list(range(start_idx, n))] + ([list(range(start_idx, -1, -1))] if start_idx > 0 else [])
    total = sum(map(len, orders))
    try:
        for order in orders:
            tracker.initialize(features, vision_pos_all, features.feat_sizes, must3r_feats, explicit, order,
                               mask_inputs=mask_inputs, point_inputs=point_inputs)
            for i in order:
                logits = tracker.step()
                done += 1
                if i not in seen:
                    seen.add(i)
                    yield i, logits_to_original_mask(logits, g), done, total
            predictor.memory_attention.clear_state()
    finally:
        predictor.memory_attention.clear_state()
        features.batches.clear()
        torch.cuda.empty_cache()


@torch.no_grad()
@torch.autocast(device_type='cuda', dtype=torch.bfloat16)
def single_frame_mask(original_predictor, frame_u8, g, points, labels, device='cuda'):
    """Original SAM2.1 prompt-to-mask on one frame (frame_u8: one sam2_frames() entry;
    points: original pixels, labels 1/0/2/3)."""
    image = sam2_batch(frame_u8[None], device)
    vf, vpe, fs = get_image_feature(original_predictor, image)
    prompts = points_to_sam2(points, labels, g, device)
    out = original_predictor.track_step(
        frame_idx=0, is_init_cond_frame=True, current_vision_feats=vf, current_vision_pos_embeds=vpe, feat_sizes=fs,
        point_inputs=prompts, mask_inputs=None, output_dict={'cond_frame_outputs': {}, 'non_cond_frame_outputs': {}},
        num_frames=1, track_in_reverse=False, run_mem_encoder=False, prev_sam_mask_logits=None)
    logits = F.interpolate(out['pred_masks'], size=(SAM2_IMAGE_SIZE, SAM2_IMAGE_SIZE), mode='bilinear', align_corners=False)
    return logits_to_original_mask(logits.squeeze().float(), g)