sam + matany2

.gitignore

@@ -17,3 +17,4 @@ utility/userdata/info/
 text_embed_cache/
 web/power_spline_editor/bg/ref_image_0.png
 web/power_spline_editor/ref/
+/web/power_spline_editor/bg

__init__.py

@@ -9,6 +9,8 @@ from .nodes.decode_overlap import *
 
 from .nodes.power_spline_editor import *
 from .nodes.prepare_refs import *
+from .nodes.video_inpaint import *
+from .nodes.matanyone2 import *
 
 from .nodes.lora_extract_v2 import *
 from .nodes.lora_change import *

api.py

@@ -8,12 +8,18 @@ from io import BytesIO
 import folder_paths
 import subprocess
 import sys
+import cv2
+import numpy as np
 
 
 import os
 import folder_paths
 from datetime import datetime
 
+# Import SAM2 masker for point-click masking
+from .nodes.sam2_masker import sam2_masker
+from .nodes.prepare_refs import convert_mask_to_contour
+
 def get_bg_folder_path():
     """Get the path to the bg folder where A.jpg and bg_image.png are stored"""
     # Look for the bg folder relative to this file
@@ -396,3 +402,133 @@ def _rename_preview_files(old_base_name, new_name, preview_folder):
 
     return renamed_files
 
+@PromptServer.instance.routes.post("/wanvideowrapper_qq/extract_first_frame")
+async def extract_first_frame(request):
+    """API endpoint to extract the first frame from a video file"""
+    try:
+        data = await request.json()
+        video_path = data.get("video_path")
+
+        if not video_path:
+            return web.json_response({"success": False, "error": "No video path provided"}, status=400)
+
+        print(f"[API extract_first_frame] Extracting first frame from: {video_path}")
+
+        # Try to find the video file
+        video_file_path = None
+
+        # Check if it's already an absolute path
+        if os.path.exists(video_path):
+            video_file_path = video_path
+        else:
+            # Search in input directories
+            input_folders = folder_paths.get_folder_paths("input")
+            for folder in input_folders:
+                potential_path = os.path.join(folder, video_path)
+                if os.path.exists(potential_path):
+                    video_file_path = potential_path
+                    break
+
+        if not video_file_path:
+            return web.json_response({"success": False, "error": f"Video file not found: {video_path}"}, status=404)
+
+        print(f"[API extract_first_frame] Found video at: {video_file_path}")
+
+        # Extract first frame using OpenCV
+        cap = cv2.VideoCapture(video_file_path)
+
+        if not cap.isOpened():
+            return web.json_response({"success": False, "error": "Failed to open video file"}, status=500)
+
+        # Read first frame
+        ret, frame = cap.read()
+        cap.release()
+
+        if not ret:
+            return web.json_response({"success": False, "error": "Failed to read first frame from video"}, status=500)
+
+        # Convert BGR to RGB
+        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+
+        # Convert to PIL Image
+        img = Image.fromarray(frame_rgb)
+
+        # Convert to base64
+        buffered = BytesIO()
+        img.save(buffered, format="PNG")
+        img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
+        img_data_url = f"data:image/png;base64,{img_base64}"
+
+        print(f"[API extract_first_frame] Successfully extracted first frame")
+
+        return web.json_response({
+            "success": True,
+            "image": img_data_url
+        })
+
+    except Exception as e:
+        print(f"[API extract_first_frame] ERROR: {e}")
+        import traceback
+        traceback.print_exc()
+        return web.json_response({"success": False, "error": str(e)}, status=500)
+
+
+@PromptServer.instance.routes.post("/wanvideowrapper_qq/sam2_predict")
+async def sam2_predict(request):
+    """
+    API endpoint for SAM2 point-click prediction.
+
+    Expects JSON with:
+        image: base64 encoded image (data URL format)
+        points: list of {'x': float, 'y': float, 'label': int} points
+                x, y are in pixel coordinates, label is 1 (foreground) or 0 (background)
+
+    Returns JSON with:
+        success: bool
+        path: list of {'x': float, 'y': float} normalized coordinates (0-1)
+        score: float (confidence score)
+        mask_shape: [height, width] of the predicted mask
+    """
+    try:
+        data = await request.json()
+
+        # Validate required parameters
+        if not data.get("image"):
+            return web.json_response({"success": False, "error": "Missing image parameter"}, status=400)
+
+        if not data.get("points") or len(data["points"]) == 0:
+            return web.json_response({"success": False, "error": "Missing points parameter"}, status=400)
+
+        # Decode base64 image
+        image_data = data["image"]
+        if image_data.startswith('data:image'):
+            header, encoded = image_data.split(',', 1)
+        else:
+            encoded = image_data
+
+        image_bytes = base64.b64decode(encoded)
+        img = Image.open(BytesIO(image_bytes))
+
+        # Extract points
+        points = data["points"]
+        point_coords = [[p["x"], p["y"]] for p in points]
+        point_labels = [p.get("label", 1) for p in points]  # Default to foreground
+
+        # Predict mask using SAM2
+        mask, score = sam2_masker.predict_from_points(img, point_coords, point_labels)
+
+        # Convert mask to contour path
+        path = convert_mask_to_contour(mask)
+
+        return web.json_response({
+            "success": True,
+            "path": path,
+            "score": score,
+            "mask_shape": mask.shape
+        })
+
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        return web.json_response({"success": False, "error": str(e)}, status=500)
+

nodes/matanyone2.py

@@ -0,0 +1,249 @@
+import os
+import sys
+import torch
+import numpy as np
+import hashlib
+import requests
+from tqdm import tqdm
+from pathlib import Path
+from typing import Optional
+
+# Add matanyone2 to path
+matanyone2_path = Path(__file__).parent.parent / "utility" / "matanyone2"
+if str(matanyone2_path) not in sys.path:
+    sys.path.insert(0, str(matanyone2_path))
+
+from matanyone.inference.inference_core import InferenceCore
+from matanyone.utils.get_default_model import get_matanyone_model
+
+
+def download_matanyone2_model(checkpoint_dir: Path = None) -> Path:
+    """Download MatAnyone2 model if not already present."""
+    MODEL_URL = "https://github.com/pq-yang/MatAnyone2/releases/download/v1.0.0/matanyone2.pth"
+    MODEL_MD5 = "b1d3cfbb7596ecf3b88391198427ca95"
+    MODEL_NAME = "matanyone2.pth"
+
+    if checkpoint_dir is None:
+        checkpoint_dir = matanyone2_path / "checkpoints"
+
+    checkpoint_dir.mkdir(parents=True, exist_ok=True)
+    model_path = checkpoint_dir / MODEL_NAME
+
+    if model_path.exists():
+        with open(model_path, 'rb') as f:
+            file_hash = hashlib.md5(f.read()).hexdigest()
+        if file_hash == MODEL_MD5:
+            return model_path
+        print(f"[MatAnyone2] Model corrupted, redownloading...")
+        model_path.unlink()
+
+    print(f"[MatAnyone2] Downloading model...")
+    try:
+        response = requests.get(MODEL_URL, stream=True)
+        response.raise_for_status()
+        total_size = int(response.headers.get('content-length', 0))
+
+        with open(model_path, 'wb') as f:
+            with tqdm(total=total_size, unit='B', unit_scale=True, desc=MODEL_NAME) as pbar:
+                for chunk in response.iter_content(chunk_size=8192):
+                    if chunk:
+                        f.write(chunk)
+                        pbar.update(len(chunk))
+
+        with open(model_path, 'rb') as f:
+            file_hash = hashlib.md5(f.read()).hexdigest()
+        if file_hash != MODEL_MD5:
+            model_path.unlink()
+            raise ValueError(f"Incorrect MD5: {file_hash}")
+
+        print(f"[MatAnyone2] Model downloaded")
+        return model_path
+
+    except Exception as e:
+        if model_path.exists():
+            model_path.unlink()
+        raise RuntimeError(f"Failed to download: {e}")
+
+
+class MatAnyone2:
+    """
+    MatAnyone2 video matting node.
+    Propagates a mask through video frames.
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "video": ("IMAGE", {"forceInput": True}),
+                "mask": ("IMAGE", {"forceInput": True}),
+            },
+            "optional": {
+                "warmup_iterations": ("INT", {"default": 5, "min": 0, "max": 10}),
+                "force_cpu": ("BOOLEAN", {"default": False}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("video",)
+    FUNCTION = "propagate_mask"
+    CATEGORY = "WanVideoWrapper_QQ"
+    DESCRIPTION = """
+    Propagate a mask through video frames using MatAnyone2.
+    video: Input video frames [B,H,W,C]
+    mask: Starting frame mask (black and white)
+    warmup_iterations: Number of refinement passes for first frame (0-10, default 5). Lower = faster but less accurate.
+    force_cpu: If False, uses GPU if available. If True, forces CPU processing.
+    """
+
+    def __init__(self):
+        self.processor = None
+
+    def _load_processor(self, device: torch.device):
+        """Load MatAnyone2 processor if not already loaded."""
+        if self.processor is not None:
+            return
+
+        checkpoint_path = download_matanyone2_model()
+
+        original_dir = os.getcwd()
+        try:
+            os.chdir(str(matanyone2_path))
+            model = get_matanyone_model(str(checkpoint_path), device=device)
+            self.processor = InferenceCore(model, cfg=model.cfg, device=device)
+        finally:
+            os.chdir(original_dir)
+
+    def _process_input_mask(self, mask: torch.Tensor) -> torch.Tensor:
+        """Convert mask to [num_objects, H, W] format for MatAnyone2.
+        Expected output: float tensor in range [0, 255] with shape [1, H, W] for single object.
+        The model will internally normalize by dividing by 255.
+        """
+        # Remove batch dimension if present: [B, H, W, C] -> [H, W, C]
+        if mask.dim() == 4:
+            mask = mask[0]
+
+        # Convert [H, W, C] to [H, W]
+        if mask.dim() == 3:
+            if mask.shape[-1] > 1:
+                # RGB/RGBA -> grayscale by averaging channels
+                mask = mask.mean(dim=-1)
+            else:
+                # Already single channel -> remove it
+                mask = mask.squeeze(-1)
+
+        # Ensure float in [0, 255] range (model will divide by 255 internally)
+        if mask.dtype == torch.uint8 or mask.dtype == torch.int8:
+            mask = mask.float()
+        elif mask.max() <= 1.0:
+            # Already in [0, 1] range -> convert to [0, 255]
+            mask = (mask * 255.0).float()
+
+        # Add object dimension: [H, W] -> [1, H, W]
+        mask = mask.unsqueeze(0)
+
+        # Ensure tensor is contiguous
+        return mask.contiguous()
+
+    def _process_frame(self, frame: torch.Tensor) -> torch.Tensor:
+        """Convert frame [H, W, C] to [3, H, W] for MatAnyone2."""
+        frame_np = (frame.cpu().numpy() * 255).astype(np.uint8)
+        if frame_np.shape[-1] == 4:
+            frame_np = frame_np[:, :, :3]
+        tensor = torch.from_numpy(frame_np / 255.0).float().permute(2, 0, 1)
+        return tensor
+
+    def _output_to_image(self, prob: torch.Tensor) -> torch.Tensor:
+        """Convert MatAnyone2 output to [H, W, C] image format.
+
+        prob can be:
+        - [H, W] - single probability map
+        - [1, H, W] - with batch/object dimension
+        - [2, H, W] - [background, foreground] probabilities
+        - [num_objects, H, W] - multiple objects
+        """
+        # Handle different input formats
+        if prob.dim() == 3 and prob.shape[0] == 2:
+            # [2, H, W] format -> take foreground (index 1)
+            prob = prob[1]  # Take foreground channel
+        elif prob.dim() == 3 and prob.shape[0] == 1:
+            # [1, H, W] format -> squeeze the batch dimension
+            prob = prob.squeeze(0)
+        elif prob.dim() > 2:
+            # Try to squeeze only dimensions of size 1
+            prob = prob.squeeze()
+
+        # Now prob should be [H, W]
+        prob_np = (prob.detach().cpu().numpy() * 255).astype(np.uint8)
+        prob_rgb = np.stack([prob_np] * 3, axis=-1)
+        return torch.from_numpy(prob_rgb / 255.0).float()
+
+    def propagate_mask(self, video: torch.Tensor, mask: torch.Tensor, warmup_iterations: int = 5, force_cpu: bool = False):
+        """
+        Propagate mask through video.
+
+        Args:
+            video: [B, H, W, C] tensor
+            mask: [B, H, W, C] or [H, W, C] tensor mask
+            warmup_iterations: Number of warmup iterations for first frame
+            force_cpu: If True, force CPU processing. If False, use GPU if available.
+
+        Returns:
+            [B, H, W, C] tensor with propagated masks
+        """
+        # Determine device: GPU by default, CPU if forced or unavailable
+        if force_cpu:
+            device = torch.device('cpu')
+        else:
+            device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+        batch_size = video.shape[0]
+
+        self._load_processor(device)
+
+        mask_input = self._process_input_mask(mask)
+
+        outputs = []
+        for i in range(batch_size):
+            frame = video[i]
+            frame_tensor = self._process_frame(frame).to(device)
+
+            # Ensure mask_input is on the correct device
+            if mask_input.device != device:
+                mask_input = mask_input.to(device)
+
+            # Ensure frame_tensor is on the correct device and has correct shape
+            if frame_tensor.device != device:
+                frame_tensor = frame_tensor.to(device)
+            if frame_tensor.dim() != 3:
+                raise ValueError(f"Expected frame_tensor to have 3 dimensions [C, H, W], got {frame_tensor.shape}")
+
+            if i == 0:
+                out_prob = self.processor.step(frame_tensor, mask_input, objects=[1])
+                for warmup_idx in range(warmup_iterations):
+                    out_prob = self.processor.step(frame_tensor, first_frame_pred=True)
+            else:
+                out_prob = self.processor.step(frame_tensor)
+
+            output_frame = self._output_to_image(out_prob.squeeze(0))
+            outputs.append(output_frame)
+
+            if device.type == 'cuda' and i % 10 == 0:
+                torch.cuda.empty_cache()
+
+        stacked = torch.stack(outputs, dim=0)
+        # Move output back to original device (usually CPU for ComfyUI)
+        output_device = video.device
+        if output_device != device:
+            stacked = stacked.to(output_device)
+
+        return (stacked,)
+
+
+NODE_CLASS_MAPPINGS = {
+    "MatAnyone2": MatAnyone2,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "MatAnyone2": "MatAnyone2",
+}

nodes/prepare_refs.py

@@ -132,13 +132,15 @@ def filter_layers_with_shapes(ref_layer_data: List[Dict[str, Any]]) -> List[Dict
 
         # optional: only include layers explicitly turned "on" if present
         if "on" in layer and not layer.get("on", False):
-            print(f"[PrepareRefs INFO] Layer {idx} ({layer.get('name','unknown')}) 'on' flag is False - skipping")
+            print(f"[PrepareRefs WARNING] Layer {idx} ({layer.get('name','unknown')}) 'on' flag is False - skipping")
             continue
 
         valid_layers.append(layer)
 
     if len(valid_layers) < len(ref_layer_data):
-        print(f"[PrepareRefs INFO] Filtered layers: {len(ref_layer_data)} → {len(valid_layers)} (removed {len(ref_layer_data)-len(valid_layers)} empty/invalid)")
+        # Filtered and valid layers processed
+        if len(valid_layers) < len(ref_layer_data):
+            pass  # Silently skip invalid layers
     return valid_layers
 
 
@@ -211,6 +213,51 @@ def create_image_and_mask_from_layer(base_image: torch.Tensor,
         return None, None
 
 
+def convert_mask_to_contour(mask_binary: np.ndarray) -> List[Dict[str, float]]:
+    """
+    Convert binary mask to normalized path points for compatibility with existing lasso system.
+
+    This function extracts the largest contour from a binary mask and converts it to
+    normalized coordinates (0-1 range) that can be used with the existing lasso system.
+
+    Args:
+        mask_binary: [H, W] numpy array (binary mask with values 0 or 1)
+
+    Returns:
+        List of {'x': float, 'y': float} in normalized 0-1 coordinates
+        Returns empty list if no contours found
+    """
+    try:
+        # Find contours
+        contours, _ = cv2.findContours(
+            mask_binary.astype(np.uint8),
+            cv2.RETR_EXTERNAL,
+            cv2.CHAIN_APPROX_SIMPLE
+        )
+
+        if not contours:
+            return []
+
+        # Get largest contour by area
+        contour = max(contours, key=cv2.contourArea)
+
+        # Simplify contour to reduce points (0.1% of contour length)
+        epsilon = 0.001 * cv2.arcLength(contour, True)
+        approx = cv2.approxPolyDP(contour, epsilon, True)
+
+        # Convert to normalized coordinates (0-1 range)
+        h, w = mask_binary.shape
+        normalized_points = [
+            {'x': float(pt[0][0]) / w, 'y': float(pt[0][1]) / h}
+            for pt in approx
+        ]
+
+        return normalized_points
+    except Exception as e:
+        print(f"[PrepareRefs ERROR] convert_mask_to_contour failed: {e}")
+        return []
+
+
 def extract_lasso_shapes_as_images(base_image: torch.Tensor,
                                    width: int,
                                    height: int,
@@ -543,8 +590,8 @@ class PrepareRefs:
             }
         }
 
-    RETURN_TYPES = ("IMAGE", "IMAGE", "MASK")
-    RETURN_NAMES = ("bg_image", "ref_images", "ref_masks")
+    RETURN_TYPES = ("IMAGE", "IMAGE", "MASK", "IMAGE")
+    RETURN_NAMES = ("bg_image", "ref_images", "ref_masks", "mat_masks")
     FUNCTION = "prepare"
     CATEGORY = "WanVideoWrapper_QQ"
     DESCRIPTION = """
@@ -597,6 +644,7 @@ class PrepareRefs:
             # Build empty outputs matching expected shapes
             empty_ref_images = torch.zeros((1, height, width, 4 if export_alpha else 3), dtype=torch.float32)
             empty_ref_masks = torch.zeros((1, height, width), dtype=torch.float32)
+            empty_mat_masks = torch.zeros((1, height, width, 3), dtype=torch.float32)
             ui_out = {"bg_image_dims": [{"width": float(width), "height": float(height)}]}
             # Persist bg preview only if original bg_image was provided
             if bg_image is not None:
@@ -605,15 +653,23 @@ class PrepareRefs:
                     ui_out["bg_image_path"] = ["bg/bg_image.png"]
                 except Exception as exc:
                     print(f"[PrepareRefs WARNING] Failed to save bg preview: {exc}")
-            return {"ui": ui_out, "result": (output_bg_image, empty_ref_images, empty_ref_masks)}
+            return {"ui": ui_out, "result": (output_bg_image, empty_ref_images, empty_ref_masks, empty_mat_masks)}
 
         # Extract per-layer images and masks (full-size)
         # Use the actual dimensions of internal_processing_image to ensure masks match
         mask_width = internal_processing_image.shape[2] if internal_processing_image is not None else width
         mask_height = internal_processing_image.shape[1] if internal_processing_image is not None else height
-        print(f"[PrepareRefs INFO] Extracting lasso shapes at {mask_width}x{mask_height} (internal_processing_image dimensions)")
         ref_images, ref_masks = extract_lasso_shapes_as_images(internal_processing_image, mask_width, mask_height, valid_ref_layers, export_alpha)
 
+        # Create mat_masks (full-size masks as RGB images for MatAnyone2)
+        # These are white-on-black masks at full resolution, not cropped to 768x768
+        if ref_masks is not None and ref_masks.shape[0] > 0:
+            # Convert masks [N, H, W] to RGB images [N, H, W, 3] (white on black)
+            mat_masks = ref_masks.unsqueeze(-1).repeat(1, 1, 1, 3)  # [N, H, W, 3]
+        else:
+            # Empty mat_masks if no masks
+            mat_masks = torch.zeros((1, height, width, 3), dtype=torch.float32)
+
         # Define fixed square size for all ref images
         FIXED_SQUARE_SIZE = 768
 
@@ -625,10 +681,12 @@ class PrepareRefs:
                 # extra_refs will be processed to 768x768, so ref_images must match for concatenation
                 ref_images = torch.zeros((0, FIXED_SQUARE_SIZE, FIXED_SQUARE_SIZE, 4 if export_alpha else 3), dtype=torch.float32)
                 ref_masks = torch.zeros((0, FIXED_SQUARE_SIZE, FIXED_SQUARE_SIZE), dtype=torch.float32)
+                mat_masks = torch.zeros((0, height, width, 3), dtype=torch.float32)
             else:
                 # No lasso layers and no extra_refs, create a single empty image
                 ref_images = torch.zeros((1, height, width, 4 if export_alpha else 3), dtype=torch.float32)
                 ref_masks = torch.zeros((1, height, width), dtype=torch.float32)
+                mat_masks = torch.zeros((1, height, width, 3), dtype=torch.float32)
 
         # Optionally save a combined mask preview and create masked BG preview
         if ref_masks is not None and ref_masks.shape[0] > 0 and to_bounding_box:
@@ -778,6 +836,33 @@ class PrepareRefs:
                 ref_images = torch.cat([ref_images, extra_refs_tensor], dim=0)
                 ref_masks = torch.cat([ref_masks, extra_masks_tensor], dim=0)
 
+                # Also add to mat_masks (convert extra_masks to full-size RGB)
+                extra_mat_masks = []
+                for i in range(num_extra_refs):
+                    if extra_masks is not None and i < extra_masks.shape[0]:
+                        # Get original mask and resize to full video size
+                        extra_mask = extra_masks[i]
+                        if extra_mask.ndim == 2:
+                            extra_mask_rgb = extra_mask.unsqueeze(-1).repeat(1, 1, 3)  # [H, W, 3]
+                        else:
+                            extra_mask_rgb = extra_mask.mean(dim=-1, keepdim=True).repeat(1, 1, 3)
+
+                        # Resize to video dimensions if needed
+                        if extra_mask_rgb.shape[0] != height or extra_mask_rgb.shape[1] != width:
+                            extra_mask_rgb = extra_mask_rgb.unsqueeze(0).permute(0, 3, 1, 2)  # [1, 3, H, W]
+                            extra_mask_rgb = F.interpolate(extra_mask_rgb, size=(height, width), mode='bilinear', align_corners=False)
+                            extra_mask_rgb = extra_mask_rgb.permute(0, 2, 3, 1).squeeze(0)  # [H, W, 3]
+
+                        extra_mat_masks.append(extra_mask_rgb)
+                    else:
+                        # Create empty full-size mask
+                        extra_mat_masks.append(torch.zeros((height, width, 3), dtype=torch.float32))
+
+                # Concatenate extra mat_masks
+                if extra_mat_masks:
+                    extra_mat_masks_tensor = torch.stack(extra_mat_masks, dim=0)
+                    mat_masks = torch.cat([mat_masks, extra_mat_masks_tensor], dim=0)
+
                 print(f"[PrepareRefs INFO] Added {num_extra_refs} extra refs. New ref_images shape: {ref_images.shape}")
 
         # Align batch sizes: if base batch size greater than refs, pad refs
@@ -792,6 +877,9 @@ class PrepareRefs:
             extra_masks = torch.zeros((extra_count, max_dim, max_dim), dtype=ref_masks.dtype, device=ref_masks.device)
             ref_masks = torch.cat([ref_masks, extra_masks], dim=0)
 
+            # Don't pad mat_masks - keep only the ref_layers, not one per video frame
+            # mat_masks should only have as many masks as there are ref_layers
+
         # Build ui_out
         ui_out = {"bg_image_dims": [{"width": float(width), "height": float(height)}]}
 
@@ -906,7 +994,7 @@ class PrepareRefs:
         except Exception as e:
             print(f"[PrepareRefs ERROR] failed to save ref masks: {e}")
 
-        return {"ui": ui_out, "result": (output_bg_image, ref_images, ref_masks)}
+        return {"ui": ui_out, "result": (output_bg_image, ref_images, ref_masks, mat_masks)}
 
     # -------------------------
     # Internal helper
@@ -928,7 +1016,8 @@ class PrepareRefs:
             original_w = bg_image.shape[2]
 
             # NEVER resize bg_image - keep original size at all costs
-            print(f"[PrepareRefs INFO] Keeping bg_image at ORIGINAL SIZE: {original_w}x{original_h} (not resizing to {final_w}x{final_h})")
+            # Keep original size
+            pass
 
             output_bg_image = bg_image
             internal_processing_image = bg_image

nodes/sam2_masker.py

@@ -0,0 +1,138 @@
+"""
+SAM2 Masker - Point-click masking using Segment Anything Model 2
+
+This module provides a singleton SAM2 masker for efficient model reuse
+and lazy loading of the SAM2 model for point-click masking operations.
+"""
+import numpy as np
+from PIL import Image
+import torch
+import os
+import sys
+import warnings
+
+# Disable torch.jit to avoid source code inspection issues in embedded Python
+os.environ['PYTORCH_JIT_DISABLE'] = '1'
+
+# Monkey-patch torch.jit.script to avoid JIT compilation issues
+original_jit_script = torch.jit.script
+def no_jit_script(fn, *args, **kwargs):
+    """Wrapper that disables JIT compilation and returns function as-is"""
+    return fn
+torch.jit.script = no_jit_script
+
+# Also patch torch.jit.trace to be safe
+original_jit_trace = torch.jit.trace
+def no_jit_trace(fn, *args, **kwargs):
+    """Wrapper that disables JIT tracing"""
+    return fn
+torch.jit.trace = no_jit_trace
+
+warnings.filterwarnings('ignore', message='torch.jit')
+
+# Import from local utility/sam directory
+current_dir = os.path.dirname(os.path.abspath(__file__))
+parent_dir = os.path.dirname(current_dir)
+sys.path.insert(0, parent_dir)
+
+from utility.sam.sam2_image_predictor import SAM2ImagePredictor
+
+
+class SAM2Masker:
+    """
+    Singleton SAM2 masker for efficient model reuse.
+
+    This class implements lazy loading of the SAM2 model to avoid
+    unnecessary GPU memory usage and startup time. The model is only
+    loaded on first use.
+    """
+    _instance = None
+    _predictor = None
+    _model_loaded = False
+
+    def __new__(cls):
+        """Implement singleton pattern"""
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def load_model(self, device='cuda'):
+        """
+        Load SAM2 model with lazy loading.
+
+        Args:
+            device: 'cuda' or 'cpu' - defaults to 'cuda'
+        """
+        if not self._model_loaded:
+            try:
+                # Try to load model with error handling for JIT compilation issues
+                import os
+                os.environ['SAM2_DISABLE_TORCH_COMPILE'] = '1'  # Disable torch.compile
+
+                self._predictor = SAM2ImagePredictor.from_pretrained(
+                    "facebook/sam2-hiera-base-plus",
+                    device=device
+                )
+                self._model_loaded = True
+            except Exception as e:
+                # Try fallback to CPU if CUDA fails
+                if device == 'cuda':
+                    try:
+                        self._predictor = SAM2ImagePredictor.from_pretrained(
+                            "facebook/sam2-hiera-base-plus",
+                            device='cpu'
+                        )
+                        self._model_loaded = True
+                    except Exception as e2:
+                        print(f"[SAM2] Model loading failed: {e}")
+                        raise
+                else:
+                    print(f"[SAM2] Model loading failed: {e}")
+                    raise
+
+    def predict_from_points(self, image, point_coords, point_labels):
+        """
+        Predict mask from point prompts.
+
+        Args:
+            image: PIL Image or numpy array [H, W, C] in RGB format
+            point_coords: List of (x, y) tuples or Nx2 numpy array in pixel coordinates
+            point_labels: List of 1 (foreground) or 0 (background)
+
+        Returns:
+            mask: [H, W] numpy array (binary mask)
+            score: float (confidence score for the best mask)
+        """
+        self.load_model()
+
+        # Convert PIL to numpy if needed
+        if isinstance(image, Image.Image):
+            image = np.array(image)
+
+        # Ensure RGB format (convert RGBA to RGB if needed)
+        if image.shape[-1] == 4:
+            image = image[:, :, :3]  # Remove alpha channel
+
+        # Set image for prediction
+        self._predictor.set_image(image)
+
+        # Convert to numpy arrays if needed
+        if isinstance(point_coords, list):
+            point_coords = np.array(point_coords)
+        if isinstance(point_labels, list):
+            point_labels = np.array(point_labels)
+
+        # Predict masks
+        masks, scores, _ = self._predictor.predict(
+            point_coords=point_coords,
+            point_labels=point_labels,
+            multimask_output=True
+        )
+
+        # Return best mask
+        best_idx = np.argmax(scores)
+        return masks[best_idx], float(scores[best_idx])
+
+
+# Global instance for use across the application
+sam2_masker = SAM2Masker()

nodes/video_inpaint.py

@@ -0,0 +1,216 @@
+import torch
+import torch.nn.functional as F
+import numpy as np
+import cv2
+from typing import Optional
+
+
+class VideoInpaint:
+    """
+    Simple video inpainting node that fills masked areas using inpainting.
+    Similar to Prepare Refs but simplified for video processing.
+    """
+
+    # Module-level constants from PrepareRefs
+    MASK_THRESHOLD = 0.01  # threshold to consider a mask pixel "active"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "video": ("IMAGE", {"forceInput": True}),
+                "video_mask": ("IMAGE", {"forceInput": True}),
+                "padding": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("video",)
+    FUNCTION = "inpaint_video"
+    CATEGORY = "WanVideoWrapper_QQ"
+    DESCRIPTION = """
+    Fill masked areas in video using inpainting.
+    video_mask: black and white image mask (same size as video frames)
+    padding: expand mask by N pixels (0 = off)
+    """
+
+    def inpaint_video(self, video: torch.Tensor, video_mask: torch.Tensor, padding: int):
+        """
+        Main entry point for video inpainting.
+
+        Args:
+            video: [B, H, W, C] tensor of video frames (float 0..1)
+            video_mask: [B, H, W, C] or [B, H, W] tensor mask (black and white)
+            padding: number of pixels to expand the mask
+
+        Returns:
+            inpainted_video: [B, H, W, C] tensor with masked areas filled
+        """
+        if video is None:
+            raise ValueError("video input is required")
+        if video_mask is None:
+            raise ValueError("video_mask input is required")
+
+        # Ensure video is in BHWC format
+        if video.dim() != 4:
+            raise ValueError(f"video must be 4D tensor [B,H,W,C], got shape {video.shape}")
+
+        batch_size, height, width, channels = video.shape
+
+        # Process video_mask to get single channel mask [B, H, W]
+        processed_mask = self._process_mask(video_mask, height, width)
+
+        # Apply padding to mask if requested
+        if padding > 0:
+            processed_mask = self._apply_padding(processed_mask, padding)
+
+        # Process each frame in the batch
+        inpainted_frames = []
+        for i in range(batch_size):
+            frame = video[i]  # [H, W, C]
+            mask = processed_mask[i]  # [H, W]
+
+            # Apply inpainting to this frame
+            inpainted_frame = self._inpaint_frame(frame, mask)
+            inpainted_frames.append(inpainted_frame)
+
+        # Stack back into batch
+        inpainted_video = torch.stack(inpainted_frames, dim=0)
+
+        return (inpainted_video,)
+
+    def _process_mask(self, video_mask: torch.Tensor, target_height: int, target_width: int) -> torch.Tensor:
+        """
+        Process video_mask to get a single channel mask tensor [B, H, W].
+        Handles both [B, H, W, C] and [B, H, W] formats.
+        """
+        # Ensure we have a batch dimension
+        if video_mask.dim() == 3:
+            # [H, W, C] or [H, W]
+            if video_mask.shape[-1] in (1, 3, 4):
+                video_mask = video_mask.unsqueeze(0)  # [1, H, W, C]
+            else:
+                video_mask = video_mask.unsqueeze(0).unsqueeze(-1)  # [1, H, W, 1]
+        elif video_mask.dim() == 2:
+            video_mask = video_mask.unsqueeze(0).unsqueeze(-1)  # [1, H, W, 1]
+
+        # Now video_mask should be [B, H, W, C] or [B, H, W, 1]
+        mask_batch = video_mask.shape[0]
+        mask_h = video_mask.shape[1]
+        mask_w = video_mask.shape[2]
+
+        # Extract single channel from mask
+        if video_mask.dim() == 4 and video_mask.shape[-1] > 1:
+            # For multi-channel masks, average the channels (grayscale)
+            processed = video_mask.mean(dim=-1)  # [B, H, W]
+        else:
+            # Squeeze the last dimension if it's 1
+            processed = video_mask.squeeze(-1)  # [B, H, W]
+
+        # Resize if needed
+        if mask_h != target_height or mask_w != target_width:
+            # Add channel dimension for interpolation
+            processed = processed.unsqueeze(1)  # [B, 1, H, W]
+            processed = F.interpolate(
+                processed,
+                size=(target_height, target_width),
+                mode='bilinear',
+                align_corners=False
+            )
+            processed = processed.squeeze(1)  # [B, H, W]
+
+        # Ensure mask is in 0..1 range
+        processed = torch.clamp(processed, 0.0, 1.0)
+
+        return processed
+
+    def _apply_padding(self, mask: torch.Tensor, padding: int) -> torch.Tensor:
+        """
+        Expand the mask by padding pixels using morphological dilation.
+
+        Args:
+            mask: [B, H, W] tensor
+            padding: number of pixels to expand
+
+        Returns:
+            padded_mask: [B, H, W] tensor with expanded mask
+        """
+        padded_masks = []
+
+        for i in range(mask.shape[0]):
+            mask_np = mask[i].cpu().numpy()
+
+            # Convert to binary (0 or 255)
+            mask_binary = (mask_np > self.MASK_THRESHOLD).astype(np.uint8) * 255
+
+            # Create a kernel for dilation
+            kernel_size = padding * 2 + 1
+            kernel = np.ones((kernel_size, kernel_size), np.uint8)
+
+            # Apply dilation to expand the mask
+            dilated = cv2.dilate(mask_binary, kernel, iterations=1)
+
+            # Convert back to float 0..1
+            dilated_tensor = torch.from_numpy(dilated.astype(np.float32) / 255.0)
+            padded_masks.append(dilated_tensor)
+
+        return torch.stack(padded_masks).to(mask.device)
+
+    def _inpaint_frame(self, frame: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        """
+        Inpaint a single frame using the mask.
+        Uses the same inpainting approach as PrepareRefs.
+
+        Args:
+            frame: [H, W, C] tensor (float 0..1)
+            mask: [H, W] tensor (float 0..1)
+
+        Returns:
+            inpainted_frame: [H, W, C] tensor
+        """
+        # Check if mask has any active areas
+        if not torch.any(mask > self.MASK_THRESHOLD):
+            return frame
+
+        # Convert to numpy for OpenCV
+        frame_np = frame.mul(255).byte().cpu().numpy()  # HWC uint8
+
+        # Ensure RGB format
+        if frame_np.shape[-1] != 3:
+            # If not RGB, try to use first 3 channels
+            if frame_np.shape[-1] >= 3:
+                frame_np = frame_np[..., :3]
+            else:
+                # Grayscale, convert to RGB
+                frame_np = cv2.cvtColor(frame_np, cv2.COLOR_GRAY2RGB)
+
+        # Prepare mask
+        mask_np = (mask.cpu().numpy() > self.MASK_THRESHOLD).astype(np.uint8) * 255
+
+        try:
+            # Convert to BGR for OpenCV
+            frame_bgr = cv2.cvtColor(frame_np, cv2.COLOR_RGB2BGR)
+
+            # Apply inpainting
+            inpainted_bgr = cv2.inpaint(frame_bgr, mask_np, 3, cv2.INPAINT_TELEA)
+
+            # Convert back to RGB
+            inpainted_rgb = cv2.cvtColor(inpainted_bgr, cv2.COLOR_BGR2RGB)
+
+            # Convert back to tensor
+            inpainted_tensor = torch.from_numpy(inpainted_rgb).float().div(255.0).to(frame.device)
+
+            return inpainted_tensor
+        except Exception as e:
+            print(f"[VideoInpaint ERROR] inpaint_frame failed: {e}")
+            return frame
+
+
+# Node registration for ComfyUI
+NODE_CLASS_MAPPINGS = {
+    "VideoInpaint": VideoInpaint,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "VideoInpaint": "Video Inpaint",
+}

utility/matanyone2/configs/efficienttam_s_512x512.yaml

@@ -0,0 +1,120 @@
+# @package _global_
+
+# Model
+model:
+  _target_: sam2.modeling.efficienttam_base.EfficientTAMBase
+  image_encoder:
+    _target_: sam2.modeling.backbones.image_encoder.ImageEncoder
+    scalp: 0
+    trunk:
+      _target_: sam2.modeling.backbones.vitdet.ViT
+      patch_size: 16
+      embed_dim: 384
+      depth: 12
+      num_heads: 6
+      mlp_ratio: 4.0
+      qkv_bias: true
+      drop_path_rate: 0.0
+      use_rel_pos: false
+      window_size: 14
+      window_block_indexes: [0, 1, 3, 4, 6, 7, 9, 10]
+    neck:
+      _target_: sam2.modeling.backbones.image_encoder.ViTDetNeck
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 256
+        normalize: true
+        scale: null
+        temperature: 10000
+      d_model: 256
+      backbone_channel_list: [384,]
+      neck_norm: LN
+
+  memory_attention:
+    _target_: sam2.modeling.memory_attention.MemoryAttention
+    d_model: 256
+    pos_enc_at_input: true
+    layer:
+      _target_: sam2.modeling.memory_attention.MemoryAttentionLayer
+      activation: relu
+      dim_feedforward: 2048
+      dropout: 0.1
+      pos_enc_at_attn: false
+      self_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [32, 32]
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+      d_model: 256
+      pos_enc_at_cross_attn_keys: true
+      pos_enc_at_cross_attn_queries: false
+      cross_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [32, 32]
+        rope_k_repeat: True
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+        kv_in_dim: 64
+    num_layers: 4
+
+  memory_encoder:
+      _target_: sam2.modeling.memory_encoder.MemoryEncoder
+      out_dim: 64
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 64
+        normalize: true
+        scale: null
+        temperature: 10000
+      mask_downsampler:
+        _target_: sam2.modeling.memory_encoder.MaskDownSampler
+        kernel_size: 3
+        stride: 2
+        padding: 1
+      fuser:
+        _target_: sam2.modeling.memory_encoder.Fuser
+        layer:
+          _target_: sam2.modeling.memory_encoder.CXBlock
+          dim: 256
+          kernel_size: 7
+          padding: 3
+          layer_scale_init_value: 1e-6
+          use_dwconv: True  # depth-wise convs
+        num_layers: 2
+
+  num_maskmem: 7
+  image_size: 512
+  # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask
+  # SAM decoder
+  sigmoid_scale_for_mem_enc: 20.0
+  sigmoid_bias_for_mem_enc: -10.0
+  use_mask_input_as_output_without_sam: true
+  # Memory
+  directly_add_no_mem_embed: true
+  use_high_res_features_in_sam: false
+  # output 3 masks on the first click on initial conditioning frames
+  multimask_output_in_sam: true
+  # SAM heads
+  iou_prediction_use_sigmoid: True
+  # cross-attend to object pointers from other frames in the ViT encoder
+  use_obj_ptrs_in_encoder: true
+  add_tpos_enc_to_obj_ptrs: false
+  only_obj_ptrs_in_the_past_for_eval: true
+  # object occlusion prediction
+  pred_obj_scores: true
+  pred_obj_scores_mlp: true
+  fixed_no_obj_ptr: true
+  # multimask tracking settings
+  multimask_output_for_tracking: true
+  use_multimask_token_for_obj_ptr: true
+  multimask_min_pt_num: 0
+  multimask_max_pt_num: 1
+  use_mlp_for_obj_ptr_proj: true
+  # Compilation flag
+  compile_image_encoder: False

utility/matanyone2/configs/eval_matanyone_config.yaml

@@ -0,0 +1,107 @@
+defaults:
+  - _self_
+  - model: base
+  - override hydra/job_logging: custom-no-rank.yaml
+
+hydra:
+  run:
+    dir: ../output/${exp_id}/${dataset}
+  output_subdir: ${now:%Y-%m-%d_%H-%M-%S}-hydra
+
+amp: False
+weights: pretrained_models/matanyone.pth # default (can be modified from outside)
+output_dir: null # defaults to run_dir; specify this to override
+flip_aug: False
+
+
+# maximum shortest side of the input; -1 means no resizing
+# With eval_vos.py, we usually just use the dataset's size (resizing done in dataloader)
+# this parameter is added for the sole purpose for the GUI in the current codebase
+# InferenceCore will downsize the input and restore the output to the original size if needed
+# if you are using this code for some other project, you can also utilize this parameter
+max_internal_size: -1
+
+# these parameters, when set, override the dataset's default; useful for debugging
+save_all: True
+use_all_masks: False
+use_long_term: False
+mem_every: 5
+
+# only relevant when long_term is not enabled
+max_mem_frames: 5
+
+# only relevant when long_term is enabled
+long_term:
+  count_usage: True
+  max_mem_frames: 10
+  min_mem_frames: 5
+  num_prototypes: 128
+  max_num_tokens: 10000
+  buffer_tokens: 2000
+
+top_k: 30
+stagger_updates: 5
+chunk_size: -1 # number of objects to process in parallel; -1 means unlimited
+save_scores: False
+save_aux: False
+visualize: False
+
+model:
+    pixel_mean: [0.485, 0.456, 0.406]
+    pixel_std: [0.229, 0.224, 0.225]
+
+    pixel_dim: 256
+    key_dim: 64
+    value_dim: 256
+    sensory_dim: 256
+    embed_dim: 256
+
+    pixel_encoder:
+      type: resnet50
+      ms_dims: [1024, 512, 256, 64, 3] # f16, f8, f4, f2, f1
+
+    mask_encoder:
+      type: resnet18
+      final_dim: 256
+
+    pixel_pe_scale: 32
+    pixel_pe_temperature: 128
+
+    object_transformer:
+      embed_dim: ${model.embed_dim}
+      ff_dim: 2048
+      num_heads: 8
+      num_blocks: 3
+      num_queries: 16
+      read_from_pixel:
+        input_norm: False
+        input_add_pe: False
+        add_pe_to_qkv: [True, True, False]
+      read_from_past:
+        add_pe_to_qkv: [True, True, False]
+      read_from_memory:
+        add_pe_to_qkv: [True, True, False]
+      read_from_query:
+        add_pe_to_qkv: [True, True, False]
+        output_norm: False
+      query_self_attention:
+        add_pe_to_qkv: [True, True, False]
+      pixel_self_attention:
+        add_pe_to_qkv: [True, True, False]
+
+    object_summarizer:
+      embed_dim: ${model.object_transformer.embed_dim}
+      num_summaries: ${model.object_transformer.num_queries}
+      add_pe: True
+
+    aux_loss:
+      sensory:
+        enabled: True
+        weight: 0.01
+      query:
+        enabled: True
+        weight: 0.01
+
+    mask_decoder:
+      # first value must equal embed_dim
+      up_dims: [256, 128, 128, 64, 16]

utility/matanyone2/configs/sam2.1_hiera_b+.yaml

@@ -0,0 +1,116 @@
+# @package _global_
+
+# Model
+model:
+  _target_: sam2.modeling.sam2_base.SAM2Base
+  image_encoder:
+    _target_: sam2.modeling.backbones.image_encoder.ImageEncoder
+    scalp: 1
+    trunk:
+      _target_: sam2.modeling.backbones.hieradet.Hiera
+      embed_dim: 112
+      num_heads: 2
+    neck:
+      _target_: sam2.modeling.backbones.image_encoder.FpnNeck
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 256
+        normalize: true
+        scale: null
+        temperature: 10000
+      d_model: 256
+      backbone_channel_list: [896, 448, 224, 112]
+      fpn_top_down_levels: [2, 3]  # output level 0 and 1 directly use the backbone features
+      fpn_interp_model: nearest
+
+  memory_attention:
+    _target_: sam2.modeling.memory_attention.MemoryAttention
+    d_model: 256
+    pos_enc_at_input: true
+    layer:
+      _target_: sam2.modeling.memory_attention.MemoryAttentionLayer
+      activation: relu
+      dim_feedforward: 2048
+      dropout: 0.1
+      pos_enc_at_attn: false
+      self_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [64, 64]
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+      d_model: 256
+      pos_enc_at_cross_attn_keys: true
+      pos_enc_at_cross_attn_queries: false
+      cross_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [64, 64]
+        rope_k_repeat: True
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+        kv_in_dim: 64
+    num_layers: 4
+
+  memory_encoder:
+      _target_: sam2.modeling.memory_encoder.MemoryEncoder
+      out_dim: 64
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 64
+        normalize: true
+        scale: null
+        temperature: 10000
+      mask_downsampler:
+        _target_: sam2.modeling.memory_encoder.MaskDownSampler
+        kernel_size: 3
+        stride: 2
+        padding: 1
+      fuser:
+        _target_: sam2.modeling.memory_encoder.Fuser
+        layer:
+          _target_: sam2.modeling.memory_encoder.CXBlock
+          dim: 256
+          kernel_size: 7
+          padding: 3
+          layer_scale_init_value: 1e-6
+          use_dwconv: True  # depth-wise convs
+        num_layers: 2
+
+  num_maskmem: 7
+  image_size: 1024
+  # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask
+  sigmoid_scale_for_mem_enc: 20.0
+  sigmoid_bias_for_mem_enc: -10.0
+  use_mask_input_as_output_without_sam: true
+  # Memory
+  directly_add_no_mem_embed: true
+  no_obj_embed_spatial: true
+  # use high-resolution feature map in the SAM mask decoder
+  use_high_res_features_in_sam: true
+  # output 3 masks on the first click on initial conditioning frames
+  multimask_output_in_sam: true
+  # SAM heads
+  iou_prediction_use_sigmoid: True
+  # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder
+  use_obj_ptrs_in_encoder: true
+  add_tpos_enc_to_obj_ptrs: true
+  proj_tpos_enc_in_obj_ptrs: true
+  use_signed_tpos_enc_to_obj_ptrs: true
+  only_obj_ptrs_in_the_past_for_eval: true
+  # object occlusion prediction
+  pred_obj_scores: true
+  pred_obj_scores_mlp: true
+  fixed_no_obj_ptr: true
+  # multimask tracking settings
+  multimask_output_for_tracking: true
+  use_multimask_token_for_obj_ptr: true
+  multimask_min_pt_num: 0
+  multimask_max_pt_num: 1
+  use_mlp_for_obj_ptr_proj: true
+  # Compilation flag
+  compile_image_encoder: False

utility/matanyone2/configs/sam2.1_hiera_l.yaml

@@ -0,0 +1,120 @@
+# @package _global_
+
+# Model
+model:
+  _target_: sam2.modeling.sam2_base.SAM2Base
+  image_encoder:
+    _target_: sam2.modeling.backbones.image_encoder.ImageEncoder
+    scalp: 1
+    trunk:
+      _target_: sam2.modeling.backbones.hieradet.Hiera
+      embed_dim: 144
+      num_heads: 2
+      stages: [2, 6, 36, 4]
+      global_att_blocks: [23, 33, 43]
+      window_pos_embed_bkg_spatial_size: [7, 7]
+      window_spec: [8, 4, 16, 8]
+    neck:
+      _target_: sam2.modeling.backbones.image_encoder.FpnNeck
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 256
+        normalize: true
+        scale: null
+        temperature: 10000
+      d_model: 256
+      backbone_channel_list: [1152, 576, 288, 144]
+      fpn_top_down_levels: [2, 3]  # output level 0 and 1 directly use the backbone features
+      fpn_interp_model: nearest
+
+  memory_attention:
+    _target_: sam2.modeling.memory_attention.MemoryAttention
+    d_model: 256
+    pos_enc_at_input: true
+    layer:
+      _target_: sam2.modeling.memory_attention.MemoryAttentionLayer
+      activation: relu
+      dim_feedforward: 2048
+      dropout: 0.1
+      pos_enc_at_attn: false
+      self_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [64, 64]
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+      d_model: 256
+      pos_enc_at_cross_attn_keys: true
+      pos_enc_at_cross_attn_queries: false
+      cross_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [64, 64]
+        rope_k_repeat: True
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+        kv_in_dim: 64
+    num_layers: 4
+
+  memory_encoder:
+      _target_: sam2.modeling.memory_encoder.MemoryEncoder
+      out_dim: 64
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 64
+        normalize: true
+        scale: null
+        temperature: 10000
+      mask_downsampler:
+        _target_: sam2.modeling.memory_encoder.MaskDownSampler
+        kernel_size: 3
+        stride: 2
+        padding: 1
+      fuser:
+        _target_: sam2.modeling.memory_encoder.Fuser
+        layer:
+          _target_: sam2.modeling.memory_encoder.CXBlock
+          dim: 256
+          kernel_size: 7
+          padding: 3
+          layer_scale_init_value: 1e-6
+          use_dwconv: True  # depth-wise convs
+        num_layers: 2
+
+  num_maskmem: 7
+  image_size: 1024
+  # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask
+  sigmoid_scale_for_mem_enc: 20.0
+  sigmoid_bias_for_mem_enc: -10.0
+  use_mask_input_as_output_without_sam: true
+  # Memory
+  directly_add_no_mem_embed: true
+  no_obj_embed_spatial: true
+  # use high-resolution feature map in the SAM mask decoder
+  use_high_res_features_in_sam: true
+  # output 3 masks on the first click on initial conditioning frames
+  multimask_output_in_sam: true
+  # SAM heads
+  iou_prediction_use_sigmoid: True
+  # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder
+  use_obj_ptrs_in_encoder: true
+  add_tpos_enc_to_obj_ptrs: true
+  proj_tpos_enc_in_obj_ptrs: true
+  use_signed_tpos_enc_to_obj_ptrs: true
+  only_obj_ptrs_in_the_past_for_eval: true
+  # object occlusion prediction
+  pred_obj_scores: true
+  pred_obj_scores_mlp: true
+  fixed_no_obj_ptr: true
+  # multimask tracking settings
+  multimask_output_for_tracking: true
+  use_multimask_token_for_obj_ptr: true
+  multimask_min_pt_num: 0
+  multimask_max_pt_num: 1
+  use_mlp_for_obj_ptr_proj: true
+  # Compilation flag
+  compile_image_encoder: False

utility/matanyone2/matanyone/LICENSE

@@ -0,0 +1,42 @@
+# S-Lab License 1.0  
+
+Copyright 2023 S-Lab
+
+Redistribution and use for non-commercial purpose in source and 
+binary forms, with or without modification, are permitted provided 
+that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright 
+   notice, this list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright 
+   notice, this list of conditions and the following disclaimer in 
+   the documentation and/or other materials provided with the 
+   distribution.
+
+3. Neither the name of the copyright holder nor the names of its 
+   contributors may be used to endorse or promote products derived 
+   from this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 
+"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT 
+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR 
+A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT 
+HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, 
+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT 
+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, 
+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY 
+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+In the event that redistribution and/or use for commercial purpose in 
+source or binary forms, with or without modification is required, 
+please contact the contributor(s) of the work.
+
+
+---
+For inquiries permission for commercial use, please consult our team: 
+Peiqing Yang (peiqingyang99@outlook.com),
+Dr. Shangchen Zhou (shangchenzhou@gmail.com),
+Prof. Chen Change Loy (ccloy@ntu.edu.sg)

utility/matanyone2/matanyone/__init__.py

@@ -0,0 +1,2 @@
+from matanyone.inference.inference_core import InferenceCore
+from matanyone.model.matanyone import MatAnyone

utility/matanyone2/matanyone/config/eval_matanyone_config.yaml

@@ -0,0 +1,47 @@
+defaults:
+  - _self_
+  - model: base
+  - override hydra/job_logging: custom-no-rank.yaml
+
+hydra:
+  run:
+    dir: ../output/${exp_id}/${dataset}
+  output_subdir: ${now:%Y-%m-%d_%H-%M-%S}-hydra
+
+amp: False
+weights: pretrained_models/matanyone.pth # default (can be modified from outside)
+output_dir: null # defaults to run_dir; specify this to override
+flip_aug: False
+
+
+# maximum shortest side of the input; -1 means no resizing
+# With eval_vos.py, we usually just use the dataset's size (resizing done in dataloader)
+# this parameter is added for the sole purpose for the GUI in the current codebase
+# InferenceCore will downsize the input and restore the output to the original size if needed
+# if you are using this code for some other project, you can also utilize this parameter
+max_internal_size: -1
+
+# these parameters, when set, override the dataset's default; useful for debugging
+save_all: True
+use_all_masks: False
+use_long_term: False
+mem_every: 5
+
+# only relevant when long_term is not enabled
+max_mem_frames: 5
+
+# only relevant when long_term is enabled
+long_term:
+  count_usage: True
+  max_mem_frames: 10
+  min_mem_frames: 5
+  num_prototypes: 128
+  max_num_tokens: 10000
+  buffer_tokens: 2000
+
+top_k: 30
+stagger_updates: 5
+chunk_size: -1 # number of objects to process in parallel; -1 means unlimited
+save_scores: False
+save_aux: False
+visualize: False

utility/matanyone2/matanyone/config/hydra/job_logging/custom-no-rank.yaml

@@ -0,0 +1,22 @@
+# python logging configuration for tasks
+version: 1
+formatters:
+  simple:
+    format: '[%(asctime)s][%(levelname)s] - %(message)s'
+    datefmt: '%Y-%m-%d %H:%M:%S'
+handlers:
+  console:
+    class: logging.StreamHandler
+    formatter: simple
+    stream: ext://sys.stdout
+  file:
+    class: logging.FileHandler
+    formatter: simple
+    # absolute file path
+    filename: ${hydra.runtime.output_dir}/${now:%Y-%m-%d_%H-%M-%S}-eval.log
+    mode: w
+root:
+  level: INFO
+  handlers: [console, file]
+
+disable_existing_loggers: false

utility/matanyone2/matanyone/config/hydra/job_logging/custom.yaml

@@ -0,0 +1,22 @@
+# python logging configuration for tasks
+version: 1
+formatters:
+  simple:
+    format: '[%(asctime)s][%(levelname)s][r${oc.env:LOCAL_RANK}] - %(message)s'
+    datefmt: '%Y-%m-%d %H:%M:%S'
+handlers:
+  console:
+    class: logging.StreamHandler
+    formatter: simple
+    stream: ext://sys.stdout
+  file:
+    class: logging.FileHandler
+    formatter: simple
+    # absolute file path
+    filename: ${hydra.runtime.output_dir}/${now:%Y-%m-%d_%H-%M-%S}-rank${oc.env:LOCAL_RANK}.log
+    mode: w
+root:
+  level: INFO
+  handlers: [console, file]
+
+disable_existing_loggers: false

utility/matanyone2/matanyone/config/model/base.yaml

@@ -0,0 +1,58 @@
+pixel_mean: [0.485, 0.456, 0.406]
+pixel_std: [0.229, 0.224, 0.225]
+
+pixel_dim: 256
+key_dim: 64
+value_dim: 256
+sensory_dim: 256
+embed_dim: 256
+
+pixel_encoder:
+  type: resnet50
+  ms_dims: [1024, 512, 256, 64, 3] # f16, f8, f4, f2, f1
+
+mask_encoder:
+  type: resnet18
+  final_dim: 256
+
+pixel_pe_scale: 32
+pixel_pe_temperature: 128
+
+object_transformer:
+  embed_dim: ${model.embed_dim}
+  ff_dim: 2048
+  num_heads: 8
+  num_blocks: 3
+  num_queries: 16
+  read_from_pixel:
+    input_norm: False
+    input_add_pe: False
+    add_pe_to_qkv: [True, True, False]
+  read_from_past:
+    add_pe_to_qkv: [True, True, False]
+  read_from_memory:
+    add_pe_to_qkv: [True, True, False]
+  read_from_query:
+    add_pe_to_qkv: [True, True, False]
+    output_norm: False
+  query_self_attention:
+    add_pe_to_qkv: [True, True, False]
+  pixel_self_attention:
+    add_pe_to_qkv: [True, True, False]
+
+object_summarizer:
+  embed_dim: ${model.object_transformer.embed_dim}
+  num_summaries: ${model.object_transformer.num_queries}
+  add_pe: True
+
+aux_loss:
+  sensory:
+    enabled: True
+    weight: 0.01
+  query:
+    enabled: True
+    weight: 0.01
+
+mask_decoder:
+  # first value must equal embed_dim
+  up_dims: [256, 128, 128, 64, 16]

utility/matanyone2/matanyone/inference/image_feature_store.py

@@ -0,0 +1,56 @@
+import warnings
+from typing import Iterable
+import torch
+from matanyone.model.matanyone import MatAnyone
+
+
+class ImageFeatureStore:
+    """
+    A cache for image features.
+    These features might be reused at different parts of the inference pipeline.
+    This class provide an interface for reusing these features.
+    It is the user's responsibility to delete redundant features.
+
+    Feature of a frame should be associated with a unique index -- typically the frame id.
+    """
+    def __init__(self, network: MatAnyone, no_warning: bool = False):
+        self.network = network
+        self._store = {}
+        self.no_warning = no_warning
+
+    def _encode_feature(self, index: int, image: torch.Tensor, last_feats=None) -> None:
+        ms_features, pix_feat = self.network.encode_image(image, last_feats=last_feats)
+        key, shrinkage, selection = self.network.transform_key(ms_features[0])
+        self._store[index] = (ms_features, pix_feat, key, shrinkage, selection)
+   
+    def get_all_features(self, images: torch.Tensor) -> (Iterable[torch.Tensor], torch.Tensor):
+        seq_length = images.shape[0]
+        ms_features, pix_feat = self.network.encode_image(images, seq_length)
+        key, shrinkage, selection = self.network.transform_key(ms_features[0])
+        for index in range(seq_length):
+            self._store[index] = ([f[index].unsqueeze(0) for f in ms_features], pix_feat[index].unsqueeze(0), key[index].unsqueeze(0), shrinkage[index].unsqueeze(0), selection[index].unsqueeze(0))
+
+    def get_features(self, index: int,
+                     image: torch.Tensor, last_feats=None) -> (Iterable[torch.Tensor], torch.Tensor):
+        if index not in self._store:
+            self._encode_feature(index, image, last_feats)
+
+        return self._store[index][:2]
+
+    def get_key(self, index: int,
+                image: torch.Tensor, last_feats=None) -> (torch.Tensor, torch.Tensor, torch.Tensor):
+        if index not in self._store:
+            self._encode_feature(index, image, last_feats)
+
+        return self._store[index][2:]
+
+    def delete(self, index: int) -> None:
+        if index in self._store:
+            del self._store[index]
+
+    def __len__(self):
+        return len(self._store)
+
+    def __del__(self):
+        if len(self._store) > 0 and not self.no_warning:
+            warnings.warn(f'Leaking {self._store.keys()} in the image feature store')

utility/matanyone2/matanyone/inference/inference_core.py

@@ -0,0 +1,550 @@
+import logging
+from omegaconf import DictConfig
+from typing import List, Optional, Iterable, Union,Tuple
+
+import os
+import cv2
+import torch
+#import imageio
+import tempfile
+import numpy as np
+from tqdm import tqdm
+from PIL import Image
+import torch.nn.functional as F
+
+from matanyone.inference.memory_manager import MemoryManager
+from matanyone.inference.object_manager import ObjectManager
+from matanyone.inference.image_feature_store import ImageFeatureStore
+from matanyone.model.matanyone import MatAnyone
+from matanyone.utils.tensor_utils import pad_divide_by, unpad, aggregate
+from matanyone.utils.inference_utils import gen_dilate, gen_erosion, read_frame_from_videos
+
+log = logging.getLogger()
+
+
+class InferenceCore:
+
+    def __init__(self,
+                 network: Union[MatAnyone,str],
+                 cfg: DictConfig = None,
+                 *,
+                 image_feature_store: ImageFeatureStore = None,
+                 device: Union[str, torch.device] = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+                 ):
+        if isinstance(network, str):
+            network = MatAnyone.from_pretrained(network)
+        network.to(device)
+        network.eval()
+        self.network = network  
+        cfg = cfg if cfg is not None else network.cfg
+        self.cfg = cfg
+        self.mem_every = cfg.mem_every
+        stagger_updates = cfg.stagger_updates
+        self.chunk_size = cfg.chunk_size
+        self.save_aux = cfg.save_aux
+        self.max_internal_size = cfg.max_internal_size
+        self.flip_aug = cfg.flip_aug
+
+        self.curr_ti = -1
+        self.last_mem_ti = 0
+        # at which time indices should we update the sensory memory
+        if stagger_updates >= self.mem_every:
+            self.stagger_ti = set(range(1, self.mem_every + 1))
+        else:
+            self.stagger_ti = set(
+                np.round(np.linspace(1, self.mem_every, stagger_updates)).astype(int))
+        self.object_manager = ObjectManager()
+        self.memory = MemoryManager(cfg=cfg, object_manager=self.object_manager)
+
+        if image_feature_store is None:
+            self.image_feature_store = ImageFeatureStore(self.network)
+        else:
+            self.image_feature_store = image_feature_store
+
+        self.last_mask = None
+        self.last_pix_feat = None
+        self.last_msk_value = None
+
+    def clear_memory(self):
+        self.curr_ti = -1
+        self.last_mem_ti = 0
+        self.memory = MemoryManager(cfg=self.cfg, object_manager=self.object_manager)
+
+    def clear_non_permanent_memory(self):
+        self.curr_ti = -1
+        self.last_mem_ti = 0
+        self.memory.clear_non_permanent_memory()
+
+    def clear_sensory_memory(self):
+        self.curr_ti = -1
+        self.last_mem_ti = 0
+        self.memory.clear_sensory_memory()
+
+    def update_config(self, cfg):
+        self.mem_every = cfg['mem_every']
+        self.memory.update_config(cfg)
+    
+    def clear_temp_mem(self):
+        self.memory.clear_work_mem()
+        # self.object_manager = ObjectManager()
+        self.memory.clear_obj_mem()
+        # self.memory.clear_sensory_memory()
+
+    def _add_memory(self,
+                    image: torch.Tensor,
+                    pix_feat: torch.Tensor,
+                    prob: torch.Tensor,
+                    key: torch.Tensor,
+                    shrinkage: torch.Tensor,
+                    selection: torch.Tensor,
+                    *,
+                    is_deep_update: bool = True,
+                    force_permanent: bool = False) -> None:
+        """
+        Memorize the given segmentation in all memory stores.
+
+        The batch dimension is 1 if flip augmentation is not used.
+        image: RGB image, (1/2)*3*H*W
+        pix_feat: from the key encoder, (1/2)*_*H*W
+        prob: (1/2)*num_objects*H*W, in [0, 1]
+        key/shrinkage/selection: for anisotropic l2, (1/2)*_*H*W
+        selection can be None if not using long-term memory
+        is_deep_update: whether to use deep update (e.g. with the mask encoder)
+        force_permanent: whether to force the memory to be permanent
+        """
+        if prob.shape[1] == 0:
+            # nothing to add
+            log.warn('Trying to add an empty object mask to memory!')
+            return
+
+        if force_permanent:
+            as_permanent = 'all'
+        else:
+            as_permanent = 'first'
+
+        self.memory.initialize_sensory_if_needed(key, self.object_manager.all_obj_ids)
+        msk_value, sensory, obj_value, _ = self.network.encode_mask(
+            image,
+            pix_feat,
+            self.memory.get_sensory(self.object_manager.all_obj_ids),
+            prob,
+            deep_update=is_deep_update,
+            chunk_size=self.chunk_size,
+            need_weights=self.save_aux)
+        self.memory.add_memory(key,
+                               shrinkage,
+                               msk_value,
+                               obj_value,
+                               self.object_manager.all_obj_ids,
+                               selection=selection,
+                               as_permanent=as_permanent)
+        self.last_mem_ti = self.curr_ti
+        if is_deep_update:
+            self.memory.update_sensory(sensory, self.object_manager.all_obj_ids)
+        self.last_msk_value = msk_value
+
+    def _segment(self,
+                 key: torch.Tensor,
+                 selection: torch.Tensor,
+                 pix_feat: torch.Tensor,
+                 ms_features: Iterable[torch.Tensor],
+                 update_sensory: bool = True) -> torch.Tensor:
+        """
+        Produce a segmentation using the given features and the memory
+
+        The batch dimension is 1 if flip augmentation is not used.
+        key/selection: for anisotropic l2: (1/2) * _ * H * W
+        pix_feat: from the key encoder, (1/2) * _ * H * W
+        ms_features: an iterable of multiscale features from the encoder, each is (1/2)*_*H*W
+                      with strides 16, 8, and 4 respectively
+        update_sensory: whether to update the sensory memory
+
+        Returns: (num_objects+1)*H*W normalized probability; the first channel is the background
+        """
+        bs = key.shape[0]
+        if self.flip_aug:
+            assert bs == 2
+        else:
+            assert bs == 1
+
+        if not self.memory.engaged:
+            log.warn('Trying to segment without any memory!')
+            return torch.zeros((1, key.shape[-2] * 16, key.shape[-1] * 16),
+                               device=key.device,
+                               dtype=key.dtype)
+        
+        uncert_output = None
+
+        if self.curr_ti == 0: # ONLY for the first frame for prediction
+            memory_readout = self.memory.read_first_frame(self.last_msk_value, pix_feat, self.last_mask, self.network, uncert_output=uncert_output)
+        else:
+            memory_readout = self.memory.read(pix_feat, key, selection, self.last_mask, self.network, uncert_output=uncert_output, last_msk_value=self.last_msk_value, ti=self.curr_ti, 
+                                              last_pix_feat=self.last_pix_feat, last_pred_mask=self.last_mask)
+        memory_readout = self.object_manager.realize_dict(memory_readout)
+
+        sensory, _, pred_prob_with_bg = self.network.segment(ms_features,
+                                                        memory_readout,
+                                                        self.memory.get_sensory(
+                                                            self.object_manager.all_obj_ids),
+                                                        chunk_size=self.chunk_size,
+                                                        update_sensory=update_sensory)
+        # remove batch dim
+        if self.flip_aug:
+            # average predictions of the non-flipped and flipped version
+            pred_prob_with_bg = (pred_prob_with_bg[0] +
+                                 torch.flip(pred_prob_with_bg[1], dims=[-1])) / 2
+        else:
+            pred_prob_with_bg = pred_prob_with_bg[0]
+        if update_sensory:
+            self.memory.update_sensory(sensory, self.object_manager.all_obj_ids)
+        return pred_prob_with_bg
+    
+    def pred_all_flow(self, images):
+        self.total_len = images.shape[0]
+        images, self.pad = pad_divide_by(images, 16)
+        images = images.unsqueeze(0)  # add the batch dimension: (1,t,c,h,w)
+        
+        self.flows_forward, self.flows_backward = self.network.pred_forward_backward_flow(images)
+
+    def encode_all_images(self, images):
+        images, self.pad = pad_divide_by(images, 16)
+        self.image_feature_store.get_all_features(images) # t c h w
+        return images
+
+    def step(self,
+             image: torch.Tensor,
+             mask: Optional[torch.Tensor] = None,
+             objects: Optional[List[int]] = None,
+             *,
+             idx_mask: bool = False,
+             end: bool = False,
+             delete_buffer: bool = True,
+             force_permanent: bool = False,
+             matting: bool = True,
+             first_frame_pred: bool = False) -> torch.Tensor:
+        """
+        Take a step with a new incoming image.
+        If there is an incoming mask with new objects, we will memorize them.
+        If there is no incoming mask, we will segment the image using the memory.
+        In both cases, we will update the memory and return a segmentation.
+
+        image: 3*H*W
+        mask: H*W (if idx mask) or len(objects)*H*W or None
+        objects: list of object ids that are valid in the mask Tensor.
+                The ids themselves do not need to be consecutive/in order, but they need to be 
+                in the same position in the list as the corresponding mask
+                in the tensor in non-idx-mask mode.
+                objects is ignored if the mask is None. 
+                If idx_mask is False and objects is None, we sequentially infer the object ids.
+        idx_mask: if True, mask is expected to contain an object id at every pixel.
+                  If False, mask should have multiple channels with each channel representing one object.
+        end: if we are at the end of the sequence, we do not need to update memory
+            if unsure just set it to False 
+        delete_buffer: whether to delete the image feature buffer after this step
+        force_permanent: the memory recorded this frame will be added to the permanent memory
+        """
+        if objects is None and mask is not None:
+            assert not idx_mask
+            objects = list(range(1, mask.shape[0] + 1))
+
+        # resize input if needed -- currently only used for the GUI
+        resize_needed = False
+        if self.max_internal_size > 0:
+            h, w = image.shape[-2:]
+            min_side = min(h, w)
+            if min_side > self.max_internal_size:
+                resize_needed = True
+                new_h = int(h / min_side * self.max_internal_size)
+                new_w = int(w / min_side * self.max_internal_size)
+                image = F.interpolate(image.unsqueeze(0),
+                                      size=(new_h, new_w),
+                                      mode='bilinear',
+                                      align_corners=False)[0]
+                if mask is not None:
+                    if idx_mask:
+                        mask = F.interpolate(mask.unsqueeze(0).unsqueeze(0).float(),
+                                             size=(new_h, new_w),
+                                             mode='nearest-exact',
+                                             align_corners=False)[0, 0].round().long()
+                    else:
+                        mask = F.interpolate(mask.unsqueeze(0),
+                                             size=(new_h, new_w),
+                                             mode='bilinear',
+                                             align_corners=False)[0]
+
+        self.curr_ti += 1
+
+        image, self.pad = pad_divide_by(image, 16) # DONE alreay for 3DCNN!!
+        image = image.unsqueeze(0)  # add the batch dimension
+        if self.flip_aug:
+            image = torch.cat([image, torch.flip(image, dims=[-1])], dim=0)
+
+        # whether to update the working memory
+        is_mem_frame = ((self.curr_ti - self.last_mem_ti >= self.mem_every) or
+                        (mask is not None)) and (not end)
+        # segment when there is no input mask or when the input mask is incomplete
+        need_segment = (mask is None) or (self.object_manager.num_obj > 0
+                                          and not self.object_manager.has_all(objects))
+        update_sensory = ((self.curr_ti - self.last_mem_ti) in self.stagger_ti) and (not end)
+
+        # reinit if it is the first frame for prediction
+        if first_frame_pred:
+            self.curr_ti = 0
+            self.last_mem_ti = 0
+            is_mem_frame = True
+            need_segment = True
+            update_sensory = True
+
+        # encoding the image
+        ms_feat, pix_feat = self.image_feature_store.get_features(self.curr_ti, image)
+        key, shrinkage, selection = self.image_feature_store.get_key(self.curr_ti, image)
+
+        # segmentation from memory if needed
+        if need_segment:
+            pred_prob_with_bg = self._segment(key,
+                                              selection,
+                                              pix_feat,
+                                              ms_feat,
+                                              update_sensory=update_sensory)
+
+        # use the input mask if provided
+        if mask is not None:
+            # inform the manager of the new objects, and get a list of temporary id
+            # temporary ids -- indicates the position of objects in the tensor
+            # (starts with 1 due to the background channel)
+            corresponding_tmp_ids, _ = self.object_manager.add_new_objects(objects)
+
+            mask, _ = pad_divide_by(mask, 16)
+            if need_segment:
+                # merge predicted mask with the incomplete input mask
+                pred_prob_no_bg = pred_prob_with_bg[1:]
+                # use the mutual exclusivity of segmentation
+                if idx_mask:
+                    pred_prob_no_bg[:, mask > 0] = 0
+                else:
+                    pred_prob_no_bg[:, mask.max(0) > 0.5] = 0
+
+                new_masks = []
+                for mask_id, tmp_id in enumerate(corresponding_tmp_ids):
+                    if idx_mask:
+                        this_mask = (mask == objects[mask_id]).type_as(pred_prob_no_bg)
+                    else:
+                        this_mask = mask[tmp_id]
+                    if tmp_id > pred_prob_no_bg.shape[0]:
+                        new_masks.append(this_mask.unsqueeze(0))
+                    else:
+                        # +1 for padding the background channel
+                        pred_prob_no_bg[tmp_id - 1] = this_mask
+                # new_masks are always in the order of tmp_id
+                mask = torch.cat([pred_prob_no_bg, *new_masks], dim=0)
+            elif idx_mask:
+                # simply convert cls to one-hot representation
+                if len(objects) == 0:
+                    if delete_buffer:
+                        self.image_feature_store.delete(self.curr_ti)
+                    log.warn('Trying to insert an empty mask as memory!')
+                    return torch.zeros((1, key.shape[-2] * 16, key.shape[-1] * 16),
+                                       device=key.device,
+                                       dtype=key.dtype)
+                mask = torch.stack(
+                    [mask == objects[mask_id] for mask_id, _ in enumerate(corresponding_tmp_ids)],
+                    dim=0)
+
+            if matting:
+                mask = mask.unsqueeze(0).float() / 255.
+                pred_prob_with_bg = torch.cat([1-mask, mask], 0)
+            else:
+                pred_prob_with_bg = aggregate(mask, dim=0)
+                pred_prob_with_bg = torch.softmax(pred_prob_with_bg, dim=0)
+
+        # Ensure last_mask always has shape [batch, num_objects, H, W]
+        # Sometimes pred_prob_with_bg[1:] is 3D [num_objects, H, W], need to add batch dim
+        self.last_mask = pred_prob_with_bg[1:]
+        if self.last_mask.dim() == 3:
+            self.last_mask = self.last_mask.unsqueeze(0)  # Add batch dimension
+        if self.flip_aug:
+            self.last_mask = torch.cat(
+                [self.last_mask, torch.flip(self.last_mask, dims=[-1])], dim=0)
+        self.last_pix_feat = pix_feat
+
+        # save as memory if needed
+        if is_mem_frame or force_permanent:
+            # clear the memory for given mask and add the first predicted mask
+            if first_frame_pred:
+                self.clear_temp_mem()
+            self._add_memory(image,
+                             pix_feat,
+                             self.last_mask,
+                             key,
+                             shrinkage,
+                             selection,
+                             force_permanent=force_permanent,
+                             is_deep_update=True)
+        else: # compute self.last_msk_value for non-memory frame 
+            msk_value, _, _, _ = self.network.encode_mask(
+                            image,
+                            pix_feat,
+                            self.memory.get_sensory(self.object_manager.all_obj_ids),
+                            self.last_mask,
+                            deep_update=False,
+                            chunk_size=self.chunk_size,
+                            need_weights=self.save_aux)
+            self.last_msk_value = msk_value
+
+        if delete_buffer:
+            self.image_feature_store.delete(self.curr_ti)
+
+        output_prob = unpad(pred_prob_with_bg, self.pad)
+        if resize_needed:
+            # restore output to the original size
+            output_prob = F.interpolate(output_prob.unsqueeze(0),
+                                        size=(h, w),
+                                        mode='bilinear',
+                                        align_corners=False)[0]
+
+        return output_prob
+
+    def delete_objects(self, objects: List[int]) -> None:
+        """
+        Delete the given objects from the memory.
+        """
+        self.object_manager.delete_objects(objects)
+        self.memory.purge_except(self.object_manager.all_obj_ids)
+
+    def output_prob_to_mask(self, output_prob: torch.Tensor, matting: bool = True) -> torch.Tensor:
+        if matting:
+            new_mask = output_prob[1:].squeeze(0)
+        else:
+            mask = torch.argmax(output_prob, dim=0)
+
+            # index in tensor != object id -- remap the ids here
+            new_mask = torch.zeros_like(mask)
+            for tmp_id, obj in self.object_manager.tmp_id_to_obj.items():
+                new_mask[mask == tmp_id] = obj.id
+
+        return new_mask
+
+    @torch.inference_mode()
+    @torch.amp.autocast("cuda")
+    def process_video(
+        self,
+        input_path: str,
+        mask_path: str,
+        output_path: str = None,
+        n_warmup: int = 10,
+        r_erode: int = 10,
+        r_dilate: int = 10,
+        suffix: str = "",
+        save_image: bool = False,
+        max_size: int = -1,
+    ) -> Tuple:
+        """
+        Process a video for object segmentation and matting.
+        This method processes a video file by performing object segmentation and matting on each frame.
+        It supports warmup frames, mask erosion/dilation, and various output options.
+        Args:
+            input_path (str): Path to the input video file
+            mask_path (str): Path to the mask image file used for initial segmentation
+            output_path (str, optional): Directory path where output files will be saved. Defaults to a temporary directory
+            n_warmup (int, optional): Number of warmup frames to use. Defaults to 10
+            r_erode (int, optional): Erosion radius for mask processing. Defaults to 10
+            r_dilate (int, optional): Dilation radius for mask processing. Defaults to 10
+            suffix (str, optional): Suffix to append to output filename. Defaults to ""
+            save_image (bool, optional): Whether to save individual frames. Defaults to False
+            max_size (int, optional): Maximum size for frame dimension. Use -1 for no limit. Defaults to -1
+        Returns:
+            Tuple[str, str]: A tuple containing:
+                - Path to the output foreground video file (str)
+                - Path to the output alpha matte video file (str)
+        Output:
+            - Saves processed video files with foreground (_fgr) and alpha matte (_pha)
+            - If save_image=True, saves individual frames in separate directories
+        """
+        output_path = output_path if output_path is not None else tempfile.TemporaryDirectory().name
+        r_erode = int(r_erode)
+        r_dilate = int(r_dilate)
+        n_warmup = int(n_warmup)
+        max_size = int(max_size)
+
+        vframes, fps, length, video_name = read_frame_from_videos(input_path)
+        repeated_frames = vframes[0].unsqueeze(0).repeat(n_warmup, 1, 1, 1)
+        vframes = torch.cat([repeated_frames, vframes], dim=0).float()
+        length += n_warmup
+
+        new_h, new_w = vframes.shape[-2:]
+        if max_size > 0:
+            h, w = new_h, new_w
+            min_side = min(h, w)
+            if min_side > max_size:
+                new_h = int(h / min_side * max_size)
+                new_w = int(w / min_side * max_size)
+                vframes = F.interpolate(vframes, size=(new_h, new_w), mode="area")
+
+        os.makedirs(output_path, exist_ok=True)
+        if suffix:
+            video_name = f"{video_name}_{suffix}"
+        if save_image:
+            os.makedirs(f"{output_path}/{video_name}", exist_ok=True)
+            os.makedirs(f"{output_path}/{video_name}/pha", exist_ok=True)
+            os.makedirs(f"{output_path}/{video_name}/fgr", exist_ok=True)
+
+        mask = np.array(Image.open(mask_path).convert("L"))
+        if r_dilate > 0:
+            mask = gen_dilate(mask, r_dilate, r_dilate)
+        if r_erode > 0:
+            mask = gen_erosion(mask, r_erode, r_erode)
+        
+        mask = torch.from_numpy(mask).cuda()
+        if max_size > 0:
+            mask = F.interpolate(
+                mask.unsqueeze(0).unsqueeze(0), size=(new_h, new_w), mode="nearest"
+            )[0, 0]
+
+        bgr = (np.array([120, 255, 155], dtype=np.float32) / 255).reshape((1, 1, 3))
+        objects = [1]
+
+        phas = []
+        fgrs = []
+        for ti in tqdm(range(length)):
+            image = vframes[ti]
+            image_np = np.array(image.permute(1, 2, 0))
+            image = (image / 255.0).cuda().float()
+
+            if ti == 0:
+                output_prob = self.step(image, mask, objects=objects)
+                output_prob = self.step(image, first_frame_pred=True)
+            else:
+                if ti <= n_warmup:
+                    output_prob = self.step(image, first_frame_pred=True)
+                else:
+                    output_prob = self.step(image)
+
+            mask = self.output_prob_to_mask(output_prob)
+            pha = mask.unsqueeze(2).cpu().numpy()
+            com_np = image_np / 255.0 * pha + bgr * (1 - pha)
+
+            if ti > (n_warmup - 1):
+                com_np = (com_np * 255).astype(np.uint8)
+                pha = (pha * 255).astype(np.uint8)
+                fgrs.append(com_np)
+                phas.append(pha)
+                if save_image:
+                    cv2.imwrite(
+                        f"{output_path}/{video_name}/pha/{str(ti - n_warmup).zfill(5)}.png",
+                        pha,
+                    )
+                    cv2.imwrite(
+                        f"{output_path}/{video_name}/fgr/{str(ti - n_warmup).zfill(5)}.png",
+                        com_np[..., [2, 1, 0]],
+                    )
+
+        fgrs = np.array(fgrs)
+        phas = np.array(phas)
+        
+        fgr_filename = f"{output_path}/{video_name}_fgr.mp4"
+        alpha_filename = f"{output_path}/{video_name}_pha.mp4"
+        
+        #imageio.mimwrite(fgr_filename, fgrs, fps=fps, quality=7)
+        #imageio.mimwrite(alpha_filename, phas, fps=fps, quality=7)
+        
+        return (fgr_filename,alpha_filename)

utility/matanyone2/matanyone/inference/kv_memory_store.py

@@ -0,0 +1,348 @@
+from typing import Dict, List, Optional, Literal
+from collections import defaultdict
+import torch
+
+
+def _add_last_dim(dictionary, key, new_value, prepend=False):
+    # append/prepend a new value to the last dimension of a tensor in a dictionary
+    # if the key does not exist, put the new value in
+    # append by default
+    if key in dictionary:
+        dictionary[key] = torch.cat([dictionary[key], new_value], -1)
+    else:
+        dictionary[key] = new_value
+
+
+class KeyValueMemoryStore:
+    """
+    Works for key/value pairs type storage
+    e.g., working and long-term memory
+    """
+    def __init__(self, save_selection: bool = False, save_usage: bool = False):
+        """
+        We store keys and values of objects that first appear in the same frame in a bucket.
+        Each bucket contains a set of object ids.
+        Each bucket is associated with a single key tensor
+            and a dictionary of value tensors indexed by object id.
+
+        The keys and values are stored as the concatenation of a permanent part and a temporary part.
+        """
+        self.save_selection = save_selection
+        self.save_usage = save_usage
+
+        self.global_bucket_id = 0  # does not reduce even if buckets are removed
+        self.buckets: Dict[int, List[int]] = {}  # indexed by bucket id
+        self.k: Dict[int, torch.Tensor] = {}  # indexed by bucket id
+        self.v: Dict[int, torch.Tensor] = {}  # indexed by object id
+
+        # indexed by bucket id; the end point of permanent memory
+        self.perm_end_pt: Dict[int, int] = defaultdict(int)
+
+        # shrinkage and selection are just like the keys
+        self.s = {}
+        if self.save_selection:
+            self.e = {}  # does not contain the permanent memory part
+
+        # usage
+        if self.save_usage:
+            self.use_cnt = {}  # indexed by bucket id, does not contain the permanent memory part
+            self.life_cnt = {}  # indexed by bucket id, does not contain the permanent memory part
+
+    def add(self,
+            key: torch.Tensor,
+            values: Dict[int, torch.Tensor],
+            shrinkage: torch.Tensor,
+            selection: torch.Tensor,
+            supposed_bucket_id: int = -1,
+            as_permanent: Literal['no', 'first', 'all'] = 'no') -> None:
+        """
+        key: (1/2)*C*N
+        values: dict of values ((1/2)*C*N), object ids are used as keys
+        shrinkage: (1/2)*1*N
+        selection: (1/2)*C*N
+
+        supposed_bucket_id: used to sync the bucket id between working and long-term memory
+        if provided, the input should all be in a single bucket indexed by this id
+        as_permanent: whether to store the input as permanent memory
+            'no': don't
+            'first': only store it as permanent memory if the bucket is empty
+            'all': always store it as permanent memory
+        """
+        bs = key.shape[0]
+        ne = key.shape[-1]
+        assert len(key.shape) == 3
+        assert len(shrinkage.shape) == 3
+        assert not self.save_selection or len(selection.shape) == 3
+        assert as_permanent in ['no', 'first', 'all']
+
+        # add the value and create new buckets if necessary
+        if supposed_bucket_id >= 0:
+            enabled_buckets = [supposed_bucket_id]
+            bucket_exist = supposed_bucket_id in self.buckets
+            for obj, value in values.items():
+                if bucket_exist:
+                    assert obj in self.v
+                    assert obj in self.buckets[supposed_bucket_id]
+                    _add_last_dim(self.v, obj, value, prepend=(as_permanent == 'all'))
+                else:
+                    assert obj not in self.v
+                    self.v[obj] = value
+            self.buckets[supposed_bucket_id] = list(values.keys())
+        else:
+            new_bucket_id = None
+            enabled_buckets = set()
+            for obj, value in values.items():
+                assert len(value.shape) == 3
+                if obj in self.v:
+                    _add_last_dim(self.v, obj, value, prepend=(as_permanent == 'all'))
+                    bucket_used = [
+                        bucket_id for bucket_id, object_ids in self.buckets.items()
+                        if obj in object_ids
+                    ]
+                    assert len(bucket_used) == 1  # each object should only be in one bucket
+                    enabled_buckets.add(bucket_used[0])
+                else:
+                    self.v[obj] = value
+                    if new_bucket_id is None:
+                        # create new bucket
+                        new_bucket_id = self.global_bucket_id
+                        self.global_bucket_id += 1
+                        self.buckets[new_bucket_id] = []
+                    # put the new object into the corresponding bucket
+                    self.buckets[new_bucket_id].append(obj)
+                    enabled_buckets.add(new_bucket_id)
+
+        # increment the permanent size if necessary
+        add_as_permanent = {}  # indexed by bucket id
+        for bucket_id in enabled_buckets:
+            add_as_permanent[bucket_id] = False
+            if as_permanent == 'all':
+                self.perm_end_pt[bucket_id] += ne
+                add_as_permanent[bucket_id] = True
+            elif as_permanent == 'first':
+                if self.perm_end_pt[bucket_id] == 0:
+                    self.perm_end_pt[bucket_id] = ne
+                    add_as_permanent[bucket_id] = True
+
+        # create new counters for usage if necessary
+        if self.save_usage and as_permanent != 'all':
+            new_count = torch.zeros((bs, ne), device=key.device, dtype=torch.float32)
+            new_life = torch.zeros((bs, ne), device=key.device, dtype=torch.float32) + 1e-7
+
+        # add the key to every bucket
+        for bucket_id in self.buckets:
+            if bucket_id not in enabled_buckets:
+                # if we are not adding new values to a bucket, we should skip it
+                continue
+
+            _add_last_dim(self.k, bucket_id, key, prepend=add_as_permanent[bucket_id])
+            _add_last_dim(self.s, bucket_id, shrinkage, prepend=add_as_permanent[bucket_id])
+            if not add_as_permanent[bucket_id]:
+                if self.save_selection:
+                    _add_last_dim(self.e, bucket_id, selection)
+                if self.save_usage:
+                    _add_last_dim(self.use_cnt, bucket_id, new_count)
+                    _add_last_dim(self.life_cnt, bucket_id, new_life)
+
+    def update_bucket_usage(self, bucket_id: int, usage: torch.Tensor) -> None:
+        # increase all life count by 1
+        # increase use of indexed elements
+        if not self.save_usage:
+            return
+
+        usage = usage[:, self.perm_end_pt[bucket_id]:]
+        if usage.shape[-1] == 0:
+            # if there is no temporary memory, we don't need to update
+            return
+        self.use_cnt[bucket_id] += usage.view_as(self.use_cnt[bucket_id])
+        self.life_cnt[bucket_id] += 1
+
+    def sieve_by_range(self, bucket_id: int, start: int, end: int, min_size: int) -> None:
+        # keep only the temporary elements *outside* of this range (with some boundary conditions)
+        # the permanent elements are ignored in this computation
+        # i.e., concat (a[:start], a[end:])
+        # bucket with size <= min_size are not modified
+
+        assert start >= 0
+        assert end <= 0
+
+        object_ids = self.buckets[bucket_id]
+        bucket_num_elements = self.k[bucket_id].shape[-1] - self.perm_end_pt[bucket_id]
+        if bucket_num_elements <= min_size:
+            return
+
+        if end == 0:
+            # negative 0 would not work as the end index!
+            # effectively make the second part an empty slice
+            end = self.k[bucket_id].shape[-1] + 1
+
+        p_size = self.perm_end_pt[bucket_id]
+        start = start + p_size
+
+        k = self.k[bucket_id]
+        s = self.s[bucket_id]
+        if self.save_selection:
+            e = self.e[bucket_id]
+        if self.save_usage:
+            use_cnt = self.use_cnt[bucket_id]
+            life_cnt = self.life_cnt[bucket_id]
+
+        self.k[bucket_id] = torch.cat([k[:, :, :start], k[:, :, end:]], -1)
+        self.s[bucket_id] = torch.cat([s[:, :, :start], s[:, :, end:]], -1)
+        if self.save_selection:
+            self.e[bucket_id] = torch.cat([e[:, :, :start - p_size], e[:, :, end:]], -1)
+        if self.save_usage:
+            self.use_cnt[bucket_id] = torch.cat([use_cnt[:, :start - p_size], use_cnt[:, end:]], -1)
+            self.life_cnt[bucket_id] = torch.cat([life_cnt[:, :start - p_size], life_cnt[:, end:]],
+                                                 -1)
+        for obj_id in object_ids:
+            v = self.v[obj_id]
+            self.v[obj_id] = torch.cat([v[:, :, :start], v[:, :, end:]], -1)
+
+    def remove_old_memory(self, bucket_id: int, max_len: int) -> None:
+        self.sieve_by_range(bucket_id, 0, -max_len, max_len)
+
+    def remove_obsolete_features(self, bucket_id: int, max_size: int) -> None:
+        # for long-term memory only
+        object_ids = self.buckets[bucket_id]
+
+        assert self.perm_end_pt[bucket_id] == 0  # permanent memory should be empty in LT memory
+
+        # normalize with life duration
+        usage = self.get_usage(bucket_id)
+        bs = usage.shape[0]
+
+        survivals = []
+
+        for bi in range(bs):
+            _, survived = torch.topk(usage[bi], k=max_size)
+            survivals.append(survived.flatten())
+            assert survived.shape[-1] == survivals[0].shape[-1]
+
+        self.k[bucket_id] = torch.stack(
+            [self.k[bucket_id][bi, :, survived] for bi, survived in enumerate(survivals)], 0)
+        self.s[bucket_id] = torch.stack(
+            [self.s[bucket_id][bi, :, survived] for bi, survived in enumerate(survivals)], 0)
+
+        if self.save_selection:
+            # Long-term memory does not store selection so this should not be needed
+            self.e[bucket_id] = torch.stack(
+                [self.e[bucket_id][bi, :, survived] for bi, survived in enumerate(survivals)], 0)
+        for obj_id in object_ids:
+            self.v[obj_id] = torch.stack(
+                [self.v[obj_id][bi, :, survived] for bi, survived in enumerate(survivals)], 0)
+
+        self.use_cnt[bucket_id] = torch.stack(
+            [self.use_cnt[bucket_id][bi, survived] for bi, survived in enumerate(survivals)], 0)
+        self.life_cnt[bucket_id] = torch.stack(
+            [self.life_cnt[bucket_id][bi, survived] for bi, survived in enumerate(survivals)], 0)
+
+    def get_usage(self, bucket_id: int) -> torch.Tensor:
+        # return normalized usage
+        if not self.save_usage:
+            raise RuntimeError('I did not count usage!')
+        else:
+            usage = self.use_cnt[bucket_id] / self.life_cnt[bucket_id]
+            return usage
+
+    def get_all_sliced(
+        self, bucket_id: int, start: int, end: int
+    ) -> (torch.Tensor, torch.Tensor, torch.Tensor, Dict[int, torch.Tensor], torch.Tensor):
+        # return k, sk, ek, value, normalized usage in order, sliced by start and end
+        # this only queries the temporary memory
+
+        assert start >= 0
+        assert end <= 0
+
+        p_size = self.perm_end_pt[bucket_id]
+        start = start + p_size
+
+        if end == 0:
+            # negative 0 would not work as the end index!
+            k = self.k[bucket_id][:, :, start:]
+            sk = self.s[bucket_id][:, :, start:]
+            ek = self.e[bucket_id][:, :, start - p_size:] if self.save_selection else None
+            value = {obj_id: self.v[obj_id][:, :, start:] for obj_id in self.buckets[bucket_id]}
+            usage = self.get_usage(bucket_id)[:, start - p_size:] if self.save_usage else None
+        else:
+            k = self.k[bucket_id][:, :, start:end]
+            sk = self.s[bucket_id][:, :, start:end]
+            ek = self.e[bucket_id][:, :, start - p_size:end] if self.save_selection else None
+            value = {obj_id: self.v[obj_id][:, :, start:end] for obj_id in self.buckets[bucket_id]}
+            usage = self.get_usage(bucket_id)[:, start - p_size:end] if self.save_usage else None
+
+        return k, sk, ek, value, usage
+
+    def purge_except(self, obj_keep_idx: List[int]):
+        # purge certain objects from the memory except the one listed
+        obj_keep_idx = set(obj_keep_idx)
+
+        # remove objects that are not in the keep list from the buckets
+        buckets_to_remove = []
+        for bucket_id, object_ids in self.buckets.items():
+            self.buckets[bucket_id] = [obj_id for obj_id in object_ids if obj_id in obj_keep_idx]
+            if len(self.buckets[bucket_id]) == 0:
+                buckets_to_remove.append(bucket_id)
+
+        # remove object values that are not in the keep list
+        self.v = {k: v for k, v in self.v.items() if k in obj_keep_idx}
+
+        # remove buckets that are empty
+        for bucket_id in buckets_to_remove:
+            del self.buckets[bucket_id]
+            del self.k[bucket_id]
+            del self.s[bucket_id]
+            if self.save_selection:
+                del self.e[bucket_id]
+            if self.save_usage:
+                del self.use_cnt[bucket_id]
+                del self.life_cnt[bucket_id]
+
+    def clear_non_permanent_memory(self):
+        # clear all non-permanent memory
+        for bucket_id in self.buckets:
+            self.sieve_by_range(bucket_id, 0, 0, 0)
+
+    def get_v_size(self, obj_id: int) -> int:
+        return self.v[obj_id].shape[-1]
+
+    def size(self, bucket_id: int) -> int:
+        if bucket_id not in self.k:
+            return 0
+        else:
+            return self.k[bucket_id].shape[-1]
+
+    def perm_size(self, bucket_id: int) -> int:
+        return self.perm_end_pt[bucket_id]
+
+    def non_perm_size(self, bucket_id: int) -> int:
+        return self.size(bucket_id) - self.perm_size(bucket_id)
+
+    def engaged(self, bucket_id: Optional[int] = None) -> bool:
+        if bucket_id is None:
+            return len(self.buckets) > 0
+        else:
+            return bucket_id in self.buckets
+
+    @property
+    def num_objects(self) -> int:
+        return len(self.v)
+
+    @property
+    def key(self) -> Dict[int, torch.Tensor]:
+        return self.k
+
+    @property
+    def value(self) -> Dict[int, torch.Tensor]:
+        return self.v
+
+    @property
+    def shrinkage(self) -> Dict[int, torch.Tensor]:
+        return self.s
+
+    @property
+    def selection(self) -> Dict[int, torch.Tensor]:
+        return self.e
+
+    def __contains__(self, key):
+        return key in self.v

utility/matanyone2/matanyone/inference/memory_manager.py

@@ -0,0 +1,455 @@
+import logging
+from omegaconf import DictConfig
+from typing import List, Dict
+import torch
+import cv2
+
+from matanyone.inference.object_manager import ObjectManager
+from matanyone.inference.kv_memory_store import KeyValueMemoryStore
+from matanyone.model.matanyone import MatAnyone
+from matanyone.model.utils.memory_utils import *
+
+log = logging.getLogger()
+
+
+class MemoryManager:
+    """
+    Manages all three memory stores and the transition between working/long-term memory
+    """
+    def __init__(self, cfg: DictConfig, object_manager: ObjectManager):
+        self.object_manager = object_manager
+        self.sensory_dim = cfg.model.sensory_dim
+        self.top_k = cfg.top_k
+        self.chunk_size = cfg.chunk_size
+
+        self.save_aux = cfg.save_aux
+
+        self.use_long_term = cfg.use_long_term
+        self.count_long_term_usage = cfg.long_term.count_usage
+        # subtract 1 because the first-frame is now counted as "permanent memory"
+        # and is not counted towards max_mem_frames
+        # but we want to keep the hyperparameters consistent as before for the same behavior
+        if self.use_long_term:
+            self.max_mem_frames = cfg.long_term.max_mem_frames - 1
+            self.min_mem_frames = cfg.long_term.min_mem_frames - 1
+            self.num_prototypes = cfg.long_term.num_prototypes
+            self.max_long_tokens = cfg.long_term.max_num_tokens
+            self.buffer_tokens = cfg.long_term.buffer_tokens
+        else:
+            self.max_mem_frames = cfg.max_mem_frames - 1
+
+        # dimensions will be inferred from input later
+        self.CK = self.CV = None
+        self.H = self.W = None
+
+        # The sensory memory is stored as a dictionary indexed by object ids
+        # each of shape bs * C^h * H * W
+        self.sensory = {}
+
+        # a dictionary indexed by object ids, each of shape bs * T * Q * C
+        self.obj_v = {}
+
+        self.work_mem = KeyValueMemoryStore(save_selection=self.use_long_term,
+                                            save_usage=self.use_long_term)
+        if self.use_long_term:
+            self.long_mem = KeyValueMemoryStore(save_usage=self.count_long_term_usage)
+
+        self.config_stale = True
+        self.engaged = False
+
+    def update_config(self, cfg: DictConfig) -> None:
+        self.config_stale = True
+        self.top_k = cfg['top_k']
+
+        assert self.use_long_term == cfg.use_long_term, 'cannot update this'
+        assert self.count_long_term_usage == cfg.long_term.count_usage, 'cannot update this'
+
+        self.use_long_term = cfg.use_long_term
+        self.count_long_term_usage = cfg.long_term.count_usage
+        if self.use_long_term:
+            self.max_mem_frames = cfg.long_term.max_mem_frames - 1
+            self.min_mem_frames = cfg.long_term.min_mem_frames - 1
+            self.num_prototypes = cfg.long_term.num_prototypes
+            self.max_long_tokens = cfg.long_term.max_num_tokens
+            self.buffer_tokens = cfg.long_term.buffer_tokens
+        else:
+            self.max_mem_frames = cfg.max_mem_frames - 1
+
+    def _readout(self, affinity, v, uncert_mask=None) -> torch.Tensor:
+        # affinity: bs*N*HW
+        # v: bs*C*N or bs*num_objects*C*N
+        # returns bs*C*HW or bs*num_objects*C*HW
+        if len(v.shape) == 3:
+            # single object
+            if uncert_mask is not None:
+                return v @ affinity * uncert_mask
+            else:
+                return v @ affinity
+        else:
+            bs, num_objects, C, N = v.shape
+            v = v.view(bs, num_objects * C, N)
+            out = v @ affinity
+            if uncert_mask is not None:
+                uncert_mask = uncert_mask.flatten(start_dim=2).expand(-1, C, -1)
+                out = out * uncert_mask
+            return out.view(bs, num_objects, C, -1)
+
+    def _get_mask_by_ids(self, mask: torch.Tensor, obj_ids: List[int]) -> torch.Tensor:
+        # -1 because the mask does not contain the background channel
+        return mask[:, [self.object_manager.find_tmp_by_id(obj) - 1 for obj in obj_ids]]
+
+    def _get_sensory_by_ids(self, obj_ids: List[int]) -> torch.Tensor:
+        return torch.stack([self.sensory[obj] for obj in obj_ids], dim=1)
+
+    def _get_object_mem_by_ids(self, obj_ids: List[int]) -> torch.Tensor:
+        return torch.stack([self.obj_v[obj] for obj in obj_ids], dim=1)
+
+    def _get_visual_values_by_ids(self, obj_ids: List[int]) -> torch.Tensor:
+        # All the values that the object ids refer to should have the same shape
+        value = torch.stack([self.work_mem.value[obj] for obj in obj_ids], dim=1)
+        if self.use_long_term and obj_ids[0] in self.long_mem.value:
+            lt_value = torch.stack([self.long_mem.value[obj] for obj in obj_ids], dim=1)
+            value = torch.cat([lt_value, value], dim=-1)
+
+        return value
+    
+    def read_first_frame(self, last_msk_value, pix_feat: torch.Tensor, 
+             last_mask: torch.Tensor, network: MatAnyone, uncert_output=None) -> Dict[int, torch.Tensor]:
+        """
+        Read from all memory stores and returns a single memory readout tensor for each object
+
+        pix_feat: (1/2) x C x H x W
+        query_key: (1/2) x C^k x H x W
+        selection:  (1/2) x C^k x H x W
+        last_mask: (1/2) x num_objects x H x W (at stride 16)
+        return a dict of memory readouts, indexed by object indices. Each readout is C*H*W
+        """
+        h, w = pix_feat.shape[-2:]
+        bs = pix_feat.shape[0]
+        assert last_mask.shape[0] == bs
+
+        """
+        Compute affinity and perform readout
+        """
+        all_readout_mem = {}
+        buckets = self.work_mem.buckets
+        for bucket_id, bucket in buckets.items():
+
+            if self.chunk_size < 1:
+                object_chunks = [bucket]
+            else:
+                object_chunks = [
+                    bucket[i:i + self.chunk_size] for i in range(0, len(bucket), self.chunk_size)
+                ]
+
+            for objects in object_chunks:
+                this_sensory = self._get_sensory_by_ids(objects)
+                this_last_mask = self._get_mask_by_ids(last_mask, objects)
+                this_msk_value = self._get_visual_values_by_ids(objects)  # (1/2)*num_objects*C*N
+                pixel_readout = network.pixel_fusion(pix_feat, last_msk_value, this_sensory,
+                                                     this_last_mask)
+                this_obj_mem = self._get_object_mem_by_ids(objects).unsqueeze(2)
+                readout_memory, aux_features = network.readout_query(pixel_readout, this_obj_mem)
+                for i, obj in enumerate(objects):
+                    all_readout_mem[obj] = readout_memory[:, i]
+
+                if self.save_aux:
+                    aux_output = {
+                        # 'sensory': this_sensory,
+                        # 'pixel_readout': pixel_readout,
+                        'q_logits': aux_features['logits'] if aux_features else None,
+                        # 'q_weights': aux_features['q_weights'] if aux_features else None,
+                        # 'p_weights': aux_features['p_weights'] if aux_features else None,
+                        # 'attn_mask': aux_features['attn_mask'].float() if aux_features else None,
+                    }
+                    self.aux = aux_output
+
+        return all_readout_mem
+
+    def read(self, pix_feat: torch.Tensor, query_key: torch.Tensor, selection: torch.Tensor,
+             last_mask: torch.Tensor, network: MatAnyone, uncert_output=None, last_msk_value=None, ti=None,
+             last_pix_feat=None, last_pred_mask=None) -> Dict[int, torch.Tensor]:
+        """
+        Read from all memory stores and returns a single memory readout tensor for each object
+
+        pix_feat: (1/2) x C x H x W
+        query_key: (1/2) x C^k x H x W
+        selection:  (1/2) x C^k x H x W
+        last_mask: (1/2) x num_objects x H x W (at stride 16)
+        return a dict of memory readouts, indexed by object indices. Each readout is C*H*W
+        """
+        h, w = pix_feat.shape[-2:]
+        bs = pix_feat.shape[0]
+        assert query_key.shape[0] == bs
+        assert selection.shape[0] == bs
+        assert last_mask.shape[0] == bs
+
+        uncert_mask = uncert_output["mask"] if uncert_output is not None else None
+
+        query_key = query_key.flatten(start_dim=2)  # bs*C^k*HW
+        selection = selection.flatten(start_dim=2)  # bs*C^k*HW
+        """
+        Compute affinity and perform readout
+        """
+        all_readout_mem = {}
+        buckets = self.work_mem.buckets
+        for bucket_id, bucket in buckets.items():
+            if self.use_long_term and self.long_mem.engaged(bucket_id):
+                # Use long-term memory
+                long_mem_size = self.long_mem.size(bucket_id)
+                memory_key = torch.cat([self.long_mem.key[bucket_id], self.work_mem.key[bucket_id]],
+                                       -1)
+                shrinkage = torch.cat(
+                    [self.long_mem.shrinkage[bucket_id], self.work_mem.shrinkage[bucket_id]], -1)
+
+                similarity = get_similarity(memory_key, shrinkage, query_key, selection)
+                affinity, usage = do_softmax(similarity,
+                                             top_k=self.top_k,
+                                             inplace=True,
+                                             return_usage=True)
+                """
+                Record memory usage for working and long-term memory
+                """
+                # ignore the index return for long-term memory
+                work_usage = usage[:, long_mem_size:]
+                self.work_mem.update_bucket_usage(bucket_id, work_usage)
+
+                if self.count_long_term_usage:
+                    # ignore the index return for working memory
+                    long_usage = usage[:, :long_mem_size]
+                    self.long_mem.update_bucket_usage(bucket_id, long_usage)
+            else:
+                # no long-term memory
+                memory_key = self.work_mem.key[bucket_id]
+                shrinkage = self.work_mem.shrinkage[bucket_id]
+                similarity = get_similarity(memory_key, shrinkage, query_key, selection, uncert_mask=uncert_mask)
+
+                if self.use_long_term:
+                    affinity, usage = do_softmax(similarity,
+                                                 top_k=self.top_k,
+                                                 inplace=True,
+                                                 return_usage=True)
+                    self.work_mem.update_bucket_usage(bucket_id, usage)
+                else:
+                    affinity = do_softmax(similarity, top_k=self.top_k, inplace=True)
+
+            if self.chunk_size < 1:
+                object_chunks = [bucket]
+            else:
+                object_chunks = [
+                    bucket[i:i + self.chunk_size] for i in range(0, len(bucket), self.chunk_size)
+                ]
+
+            for objects in object_chunks:
+                this_sensory = self._get_sensory_by_ids(objects)
+                this_last_mask = self._get_mask_by_ids(last_mask, objects)
+                this_msk_value = self._get_visual_values_by_ids(objects)  # (1/2)*num_objects*C*N
+                visual_readout = self._readout(affinity,
+                                               this_msk_value, uncert_mask).view(bs, len(objects), self.CV, h, w)
+                
+                uncert_output = network.pred_uncertainty(last_pix_feat, pix_feat, last_pred_mask, visual_readout[:,0]-last_msk_value[:,0])
+
+                if uncert_output is not None:
+                    uncert_prob = uncert_output["prob"].unsqueeze(1) # b n 1 h w
+                    visual_readout = visual_readout*uncert_prob + last_msk_value*(1-uncert_prob)
+
+                pixel_readout = network.pixel_fusion(pix_feat, visual_readout, this_sensory,
+                                                     this_last_mask)
+                this_obj_mem = self._get_object_mem_by_ids(objects).unsqueeze(2)
+                readout_memory, aux_features = network.readout_query(pixel_readout, this_obj_mem)
+                for i, obj in enumerate(objects):
+                    all_readout_mem[obj] = readout_memory[:, i]
+
+                if self.save_aux:
+                    aux_output = {
+                        # 'sensory': this_sensory,
+                        # 'pixel_readout': pixel_readout,
+                        'q_logits': aux_features['logits'] if aux_features else None,
+                        # 'q_weights': aux_features['q_weights'] if aux_features else None,
+                        # 'p_weights': aux_features['p_weights'] if aux_features else None,
+                        # 'attn_mask': aux_features['attn_mask'].float() if aux_features else None,
+                    }
+                    self.aux = aux_output
+
+        return all_readout_mem
+
+    def add_memory(self,
+                   key: torch.Tensor,
+                   shrinkage: torch.Tensor,
+                   msk_value: torch.Tensor,
+                   obj_value: torch.Tensor,
+                   objects: List[int],
+                   selection: torch.Tensor = None,
+                   *,
+                   as_permanent: bool = False) -> None:
+        # key: (1/2)*C*H*W
+        # msk_value: (1/2)*num_objects*C*H*W
+        # obj_value: (1/2)*num_objects*Q*C
+        # objects contains a list of object ids corresponding to the objects in msk_value/obj_value
+        bs = key.shape[0]
+        assert shrinkage.shape[0] == bs
+        assert msk_value.shape[0] == bs
+        assert obj_value.shape[0] == bs
+
+        self.engaged = True
+        if self.H is None or self.config_stale:
+            self.config_stale = False
+            self.H, self.W = msk_value.shape[-2:]
+            self.HW = self.H * self.W
+            # convert from num. frames to num. tokens
+            self.max_work_tokens = self.max_mem_frames * self.HW
+            if self.use_long_term:
+                self.min_work_tokens = self.min_mem_frames * self.HW
+
+        # key:   bs*C*N
+        # value: bs*num_objects*C*N
+        key = key.flatten(start_dim=2)
+        shrinkage = shrinkage.flatten(start_dim=2)
+        self.CK = key.shape[1]
+
+        msk_value = msk_value.flatten(start_dim=3)
+        self.CV = msk_value.shape[2]
+
+        if selection is not None:
+            # not used in non-long-term mode
+            selection = selection.flatten(start_dim=2)
+
+        # insert object values into object memory
+        for obj_id, obj in enumerate(objects):
+            if obj in self.obj_v:
+                """streaming average
+                each self.obj_v[obj] is (1/2)*num_summaries*(embed_dim+1)
+                first embed_dim keeps track of the sum of embeddings
+                the last dim keeps the total count
+                averaging in done inside the object transformer
+
+                incoming obj_value is (1/2)*num_objects*num_summaries*(embed_dim+1)
+                self.obj_v[obj] = torch.cat([self.obj_v[obj], obj_value[:, obj_id]], dim=0)
+                """
+                last_acc = self.obj_v[obj][:, :, -1]
+                new_acc = last_acc + obj_value[:, obj_id, :, -1]
+
+                self.obj_v[obj][:, :, :-1] = (self.obj_v[obj][:, :, :-1] +
+                                              obj_value[:, obj_id, :, :-1])
+                self.obj_v[obj][:, :, -1] = new_acc
+            else:
+                self.obj_v[obj] = obj_value[:, obj_id]
+
+        # convert mask value tensor into a dict for insertion
+        msk_values = {obj: msk_value[:, obj_id] for obj_id, obj in enumerate(objects)}
+        self.work_mem.add(key,
+                          msk_values,
+                          shrinkage,
+                          selection=selection,
+                          as_permanent=as_permanent)
+
+        for bucket_id in self.work_mem.buckets.keys():
+            # long-term memory cleanup
+            if self.use_long_term:
+                # Do memory compressed if needed
+                if self.work_mem.non_perm_size(bucket_id) >= self.max_work_tokens:
+                    # Remove obsolete features if needed
+                    if self.long_mem.non_perm_size(bucket_id) >= (self.max_long_tokens -
+                                                         self.num_prototypes):
+                        self.long_mem.remove_obsolete_features(
+                            bucket_id,
+                            self.max_long_tokens - self.num_prototypes - self.buffer_tokens)
+
+                    self.compress_features(bucket_id)
+            else:
+                # FIFO
+                self.work_mem.remove_old_memory(bucket_id, self.max_work_tokens)
+
+    def purge_except(self, obj_keep_idx: List[int]) -> None:
+        # purge certain objects from the memory except the one listed
+        self.work_mem.purge_except(obj_keep_idx)
+        if self.use_long_term and self.long_mem.engaged():
+            self.long_mem.purge_except(obj_keep_idx)
+        self.sensory = {k: v for k, v in self.sensory.items() if k in obj_keep_idx}
+
+        if not self.work_mem.engaged():
+            # everything is removed!
+            self.engaged = False
+
+    def compress_features(self, bucket_id: int) -> None:
+        HW = self.HW
+
+        # perform memory consolidation
+        prototype_key, prototype_value, prototype_shrinkage = self.consolidation(
+            *self.work_mem.get_all_sliced(bucket_id, 0, -self.min_work_tokens))
+
+        # remove consolidated working memory
+        self.work_mem.sieve_by_range(bucket_id,
+                                     0,
+                                     -self.min_work_tokens,
+                                     min_size=self.min_work_tokens)
+
+        # add to long-term memory
+        self.long_mem.add(prototype_key,
+                          prototype_value,
+                          prototype_shrinkage,
+                          selection=None,
+                          supposed_bucket_id=bucket_id)
+
+    def consolidation(self, candidate_key: torch.Tensor, candidate_shrinkage: torch.Tensor,
+                      candidate_selection: torch.Tensor, candidate_value: Dict[int, torch.Tensor],
+                      usage: torch.Tensor) -> (torch.Tensor, Dict[int, torch.Tensor], torch.Tensor):
+        # find the indices with max usage
+        bs = candidate_key.shape[0]
+        assert bs in [1, 2]
+
+        prototype_key = []
+        prototype_selection = []
+        for bi in range(bs):
+            _, max_usage_indices = torch.topk(usage[bi], k=self.num_prototypes, dim=-1, sorted=True)
+            prototype_indices = max_usage_indices.flatten()
+            prototype_key.append(candidate_key[bi, :, prototype_indices])
+            prototype_selection.append(candidate_selection[bi, :, prototype_indices])
+        prototype_key = torch.stack(prototype_key, dim=0)
+        prototype_selection = torch.stack(prototype_selection, dim=0)
+        """
+        Potentiation step
+        """
+        similarity = get_similarity(candidate_key, candidate_shrinkage, prototype_key,
+                                    prototype_selection)
+        affinity = do_softmax(similarity)
+
+        # readout the values
+        prototype_value = {k: self._readout(affinity, v) for k, v in candidate_value.items()}
+
+        # readout the shrinkage term
+        prototype_shrinkage = self._readout(affinity, candidate_shrinkage)
+
+        return prototype_key, prototype_value, prototype_shrinkage
+
+    def initialize_sensory_if_needed(self, sample_key: torch.Tensor, ids: List[int]):
+        for obj in ids:
+            if obj not in self.sensory:
+                # also initializes the sensory memory
+                bs, _, h, w = sample_key.shape
+                self.sensory[obj] = torch.zeros((bs, self.sensory_dim, h, w),
+                                                device=sample_key.device)
+
+    def update_sensory(self, sensory: torch.Tensor, ids: List[int]):
+        # sensory: 1*num_objects*C*H*W
+        for obj_id, obj in enumerate(ids):
+            self.sensory[obj] = sensory[:, obj_id]
+
+    def get_sensory(self, ids: List[int]):
+        # returns (1/2)*num_objects*C*H*W
+        return self._get_sensory_by_ids(ids)
+    
+    def clear_non_permanent_memory(self):
+        self.work_mem.clear_non_permanent_memory()
+        if self.use_long_term:
+            self.long_mem.clear_non_permanent_memory()
+
+    def clear_sensory_memory(self):
+        self.sensory = {}
+    
+    def clear_work_mem(self):
+        self.work_mem = KeyValueMemoryStore(save_selection=self.use_long_term,
+                                            save_usage=self.use_long_term)
+    
+    def clear_obj_mem(self):
+        self.obj_v = {}

utility/matanyone2/matanyone/inference/object_info.py

@@ -0,0 +1,24 @@
+class ObjectInfo:
+    """
+    Store meta information for an object
+    """
+    def __init__(self, id: int):
+        self.id = id
+        self.poke_count = 0  # count number of detections missed
+
+    def poke(self) -> None:
+        self.poke_count += 1
+
+    def unpoke(self) -> None:
+        self.poke_count = 0
+
+    def __hash__(self):
+        return hash(self.id)
+
+    def __eq__(self, other):
+        if type(other) == int:
+            return self.id == other
+        return self.id == other.id
+
+    def __repr__(self):
+        return f'(ID: {self.id})'

utility/matanyone2/matanyone/inference/object_manager.py

@@ -0,0 +1,149 @@
+from typing import Union, List, Dict
+
+import torch
+from matanyone.inference.object_info import ObjectInfo
+
+
+class ObjectManager:
+    """
+    Object IDs are immutable. The same ID always represent the same object.
+    Temporary IDs are the positions of each object in the tensor. It changes as objects get removed.
+    Temporary IDs start from 1.
+    """
+
+    def __init__(self):
+        self.obj_to_tmp_id: Dict[ObjectInfo, int] = {}
+        self.tmp_id_to_obj: Dict[int, ObjectInfo] = {}
+        self.obj_id_to_obj: Dict[int, ObjectInfo] = {}
+
+        self.all_historical_object_ids: List[int] = []
+
+    def _recompute_obj_id_to_obj_mapping(self) -> None:
+        self.obj_id_to_obj = {obj.id: obj for obj in self.obj_to_tmp_id}
+
+    def add_new_objects(
+            self, objects: Union[List[ObjectInfo], ObjectInfo,
+                                 List[int]]) -> (List[int], List[int]):
+        if not isinstance(objects, list):
+            objects = [objects]
+
+        corresponding_tmp_ids = []
+        corresponding_obj_ids = []
+        for obj in objects:
+            if isinstance(obj, int):
+                obj = ObjectInfo(id=obj)
+
+            if obj in self.obj_to_tmp_id:
+                # old object
+                corresponding_tmp_ids.append(self.obj_to_tmp_id[obj])
+                corresponding_obj_ids.append(obj.id)
+            else:
+                # new object
+                new_obj = ObjectInfo(id=obj.id)
+
+                # new object
+                new_tmp_id = len(self.obj_to_tmp_id) + 1
+                self.obj_to_tmp_id[new_obj] = new_tmp_id
+                self.tmp_id_to_obj[new_tmp_id] = new_obj
+                self.all_historical_object_ids.append(new_obj.id)
+                corresponding_tmp_ids.append(new_tmp_id)
+                corresponding_obj_ids.append(new_obj.id)
+
+        self._recompute_obj_id_to_obj_mapping()
+        assert corresponding_tmp_ids == sorted(corresponding_tmp_ids)
+        return corresponding_tmp_ids, corresponding_obj_ids
+
+    def delete_objects(self, obj_ids_to_remove: Union[int, List[int]]) -> None:
+        # delete an object or a list of objects
+        # re-sort the tmp ids
+        if isinstance(obj_ids_to_remove, int):
+            obj_ids_to_remove = [obj_ids_to_remove]
+
+        new_tmp_id = 1
+        total_num_id = len(self.obj_to_tmp_id)
+
+        local_obj_to_tmp_id = {}
+        local_tmp_to_obj_id = {}
+
+        for tmp_iter in range(1, total_num_id + 1):
+            obj = self.tmp_id_to_obj[tmp_iter]
+            if obj.id not in obj_ids_to_remove:
+                local_obj_to_tmp_id[obj] = new_tmp_id
+                local_tmp_to_obj_id[new_tmp_id] = obj
+                new_tmp_id += 1
+
+        self.obj_to_tmp_id = local_obj_to_tmp_id
+        self.tmp_id_to_obj = local_tmp_to_obj_id
+        self._recompute_obj_id_to_obj_mapping()
+
+    def purge_inactive_objects(self,
+                               max_missed_detection_count: int) -> (bool, List[int], List[int]):
+        # remove tmp ids of objects that are removed
+        obj_id_to_be_deleted = []
+        tmp_id_to_be_deleted = []
+        tmp_id_to_keep = []
+        obj_id_to_keep = []
+
+        for obj in self.obj_to_tmp_id:
+            if obj.poke_count > max_missed_detection_count:
+                obj_id_to_be_deleted.append(obj.id)
+                tmp_id_to_be_deleted.append(self.obj_to_tmp_id[obj])
+            else:
+                tmp_id_to_keep.append(self.obj_to_tmp_id[obj])
+                obj_id_to_keep.append(obj.id)
+
+        purge_activated = len(obj_id_to_be_deleted) > 0
+        if purge_activated:
+            self.delete_objects(obj_id_to_be_deleted)
+        return purge_activated, tmp_id_to_keep, obj_id_to_keep
+
+    def tmp_to_obj_cls(self, mask) -> torch.Tensor:
+        # remap tmp id cls representation to the true object id representation
+        new_mask = torch.zeros_like(mask)
+        for tmp_id, obj in self.tmp_id_to_obj.items():
+            new_mask[mask == tmp_id] = obj.id
+        return new_mask
+
+    def get_tmp_to_obj_mapping(self) -> Dict[int, ObjectInfo]:
+        # returns the mapping in a dict format for saving it with pickle
+        return {obj.id: tmp_id for obj, tmp_id in self.tmp_id_to_obj.items()}
+
+    def realize_dict(self, obj_dict, dim=1) -> torch.Tensor:
+        # turns a dict indexed by obj id into a tensor, ordered by tmp IDs
+        output = []
+        for _, obj in self.tmp_id_to_obj.items():
+            if obj.id not in obj_dict:
+                raise NotImplementedError
+            output.append(obj_dict[obj.id])
+        output = torch.stack(output, dim=dim)
+        return output
+
+    def make_one_hot(self, cls_mask) -> torch.Tensor:
+        output = []
+        for _, obj in self.tmp_id_to_obj.items():
+            output.append(cls_mask == obj.id)
+        if len(output) == 0:
+            output = torch.zeros((0, *cls_mask.shape), dtype=torch.bool, device=cls_mask.device)
+        else:
+            output = torch.stack(output, dim=0)
+        return output
+
+    @property
+    def all_obj_ids(self) -> List[int]:
+        return [k.id for k in self.obj_to_tmp_id]
+
+    @property
+    def num_obj(self) -> int:
+        return len(self.obj_to_tmp_id)
+
+    def has_all(self, objects: List[int]) -> bool:
+        for obj in objects:
+            if obj not in self.obj_to_tmp_id:
+                return False
+        return True
+
+    def find_object_by_id(self, obj_id) -> ObjectInfo:
+        return self.obj_id_to_obj[obj_id]
+
+    def find_tmp_by_id(self, obj_id) -> int:
+        return self.obj_to_tmp_id[self.obj_id_to_obj[obj_id]]

utility/matanyone2/matanyone/inference/utils/args_utils.py

@@ -0,0 +1,30 @@
+import logging
+from omegaconf import DictConfig
+
+log = logging.getLogger()
+
+
+def get_dataset_cfg(cfg: DictConfig):
+    dataset_name = cfg.dataset
+    data_cfg = cfg.datasets[dataset_name]
+
+    potential_overrides = [
+        'image_directory',
+        'mask_directory',
+        'json_directory',
+        'size',
+        'save_all',
+        'use_all_masks',
+        'use_long_term',
+        'mem_every',
+    ]
+
+    for override in potential_overrides:
+        if cfg[override] is not None:
+            log.info(f'Overriding config {override} from {data_cfg[override]} to {cfg[override]}')
+            data_cfg[override] = cfg[override]
+        # escalte all potential overrides to the top-level config
+        if override in data_cfg:
+            cfg[override] = data_cfg[override]
+
+    return data_cfg

utility/matanyone2/matanyone/model/aux_modules.py

@@ -0,0 +1,93 @@
+"""
+For computing auxiliary outputs for auxiliary losses
+"""
+from typing import Dict
+from omegaconf import DictConfig
+import torch
+import torch.nn as nn
+
+from matanyone.model.group_modules import GConv2d
+from matanyone.utils.tensor_utils import aggregate
+
+
+class LinearPredictor(nn.Module):
+    def __init__(self, x_dim: int, pix_dim: int):
+        super().__init__()
+        self.projection = GConv2d(x_dim, pix_dim + 1, kernel_size=1)
+
+    def forward(self, pix_feat: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
+        # pixel_feat: B*pix_dim*H*W
+        # x: B*num_objects*x_dim*H*W
+        num_objects = x.shape[1]
+        x = self.projection(x)
+
+        pix_feat = pix_feat.unsqueeze(1).expand(-1, num_objects, -1, -1, -1)
+        logits = (pix_feat * x[:, :, :-1]).sum(dim=2) + x[:, :, -1]
+        return logits
+
+
+class DirectPredictor(nn.Module):
+    def __init__(self, x_dim: int):
+        super().__init__()
+        self.projection = GConv2d(x_dim, 1, kernel_size=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: B*num_objects*x_dim*H*W
+        logits = self.projection(x).squeeze(2)
+        return logits
+
+
+class AuxComputer(nn.Module):
+    def __init__(self, cfg: DictConfig):
+        super().__init__()
+
+        use_sensory_aux = cfg.model.aux_loss.sensory.enabled
+        self.use_query_aux = cfg.model.aux_loss.query.enabled
+        self.use_sensory_aux = use_sensory_aux
+
+        sensory_dim = cfg.model.sensory_dim
+        embed_dim = cfg.model.embed_dim
+
+        if use_sensory_aux:  
+            self.sensory_aux = LinearPredictor(sensory_dim, embed_dim)
+
+    def _aggregate_with_selector(self, logits: torch.Tensor, selector: torch.Tensor) -> torch.Tensor:
+        prob = torch.sigmoid(logits)
+        if selector is not None:
+            prob = prob * selector
+        logits = aggregate(prob, dim=1)
+        return logits
+
+    def forward(self, pix_feat: torch.Tensor, aux_input: Dict[str, torch.Tensor],
+                selector: torch.Tensor, seg_pass=False) -> Dict[str, torch.Tensor]:
+        sensory = aux_input['sensory']
+        q_logits = aux_input['q_logits']
+
+        aux_output = {}
+        aux_output['attn_mask'] = aux_input['attn_mask']
+
+        if self.use_sensory_aux:
+            # B*num_objects*H*W
+            logits = self.sensory_aux(pix_feat, sensory)
+            aux_output['sensory_logits'] = self._aggregate_with_selector(logits, selector)
+        if self.use_query_aux:
+            # B*num_objects*num_levels*H*W
+            aux_output['q_logits'] = self._aggregate_with_selector(
+                torch.stack(q_logits, dim=2),
+                selector.unsqueeze(2) if selector is not None else None)
+
+        return aux_output
+    
+    def compute_mask(self, aux_input: Dict[str, torch.Tensor],
+                selector: torch.Tensor) -> Dict[str, torch.Tensor]:
+        # sensory = aux_input['sensory']
+        q_logits = aux_input['q_logits']
+
+        aux_output = {}
+
+        # B*num_objects*num_levels*H*W
+        aux_output['q_logits'] = self._aggregate_with_selector(
+            torch.stack(q_logits, dim=2),
+            selector.unsqueeze(2) if selector is not None else None)
+
+        return aux_output

utility/matanyone2/matanyone/model/big_modules.py

@@ -0,0 +1,365 @@
+"""
+big_modules.py - This file stores higher-level network blocks.
+
+x - usually denotes features that are shared between objects.
+g - usually denotes features that are not shared between objects 
+    with an extra "num_objects" dimension (batch_size * num_objects * num_channels * H * W).
+
+The trailing number of a variable usually denotes the stride
+"""
+
+from typing import Iterable
+from omegaconf import DictConfig
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from matanyone.model.group_modules import MainToGroupDistributor, GroupFeatureFusionBlock, GConv2d
+from matanyone.model.utils import resnet
+from matanyone.model.modules import SensoryDeepUpdater, SensoryUpdater_fullscale, DecoderFeatureProcessor, MaskUpsampleBlock
+
+class UncertPred(nn.Module):
+    def __init__(self, model_cfg: DictConfig):
+        super().__init__()
+        self.conv1x1_v2 = nn.Conv2d(model_cfg.pixel_dim*2 + 1 + model_cfg.value_dim, 64, kernel_size=1, stride=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv3x3 = nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=1, groups=1, bias=False, dilation=1)
+        self.bn2 = nn.BatchNorm2d(32)
+        self.conv3x3_out = nn.Conv2d(32, 1, kernel_size=3, stride=1, padding=1, groups=1, bias=False, dilation=1)
+    
+    def forward(self, last_frame_feat: torch.Tensor, cur_frame_feat: torch.Tensor, last_mask: torch.Tensor, mem_val_diff:torch.Tensor):
+        last_mask = F.interpolate(last_mask, size=last_frame_feat.shape[-2:], mode='area')
+        x = torch.cat([last_frame_feat, cur_frame_feat, last_mask, mem_val_diff], dim=1)
+        x = self.conv1x1_v2(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.conv3x3(x)
+        x = self.bn2(x)
+        x = self.relu(x)
+        x = self.conv3x3_out(x)
+        return x
+    
+    # override the default train() to freeze BN statistics
+    def train(self, mode=True):
+        self.training = False
+        for module in self.children():
+            module.train(False)
+        return self
+
+class PixelEncoder(nn.Module):
+    def __init__(self, model_cfg: DictConfig):
+        super().__init__()
+
+        self.is_resnet = 'resnet' in model_cfg.pixel_encoder.type
+        # if model_cfg.pretrained_resnet is set in the model_cfg we get the value
+        # else default to True
+        is_pretrained_resnet = getattr(model_cfg,"pretrained_resnet",True)
+        if self.is_resnet:
+            if model_cfg.pixel_encoder.type == 'resnet18':
+                network = resnet.resnet18(pretrained=is_pretrained_resnet)
+            elif model_cfg.pixel_encoder.type == 'resnet50':
+                network = resnet.resnet50(pretrained=is_pretrained_resnet)
+            else:
+                raise NotImplementedError
+            self.conv1 = network.conv1
+            self.bn1 = network.bn1
+            self.relu = network.relu
+            self.maxpool = network.maxpool
+
+            self.res2 = network.layer1
+            self.layer2 = network.layer2
+            self.layer3 = network.layer3
+        else:
+            raise NotImplementedError
+
+    def forward(self, x: torch.Tensor, seq_length=None) -> (torch.Tensor, torch.Tensor, torch.Tensor):
+        f1 = x
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        f2 = x
+        x = self.maxpool(x)
+        f4 = self.res2(x)
+        f8 = self.layer2(f4)
+        f16 = self.layer3(f8)
+
+        return f16, f8, f4, f2, f1
+
+    # override the default train() to freeze BN statistics
+    def train(self, mode=True):
+        self.training = False
+        for module in self.children():
+            module.train(False)
+        return self
+
+
+class KeyProjection(nn.Module):
+    def __init__(self, model_cfg: DictConfig):
+        super().__init__()
+        in_dim = model_cfg.pixel_encoder.ms_dims[0]
+        mid_dim = model_cfg.pixel_dim
+        key_dim = model_cfg.key_dim
+
+        self.pix_feat_proj = nn.Conv2d(in_dim, mid_dim, kernel_size=1)
+        self.key_proj = nn.Conv2d(mid_dim, key_dim, kernel_size=3, padding=1)
+        # shrinkage
+        self.d_proj = nn.Conv2d(mid_dim, 1, kernel_size=3, padding=1)
+        # selection
+        self.e_proj = nn.Conv2d(mid_dim, key_dim, kernel_size=3, padding=1)
+
+        nn.init.orthogonal_(self.key_proj.weight.data)
+        nn.init.zeros_(self.key_proj.bias.data)
+
+    def forward(self, x: torch.Tensor, *, need_s: bool,
+                need_e: bool) -> (torch.Tensor, torch.Tensor, torch.Tensor):
+        x = self.pix_feat_proj(x)
+        shrinkage = self.d_proj(x)**2 + 1 if (need_s) else None
+        selection = torch.sigmoid(self.e_proj(x)) if (need_e) else None
+
+        return self.key_proj(x), shrinkage, selection
+
+
+class MaskEncoder(nn.Module):
+    def __init__(self, model_cfg: DictConfig, single_object=False):
+        super().__init__()
+        pixel_dim = model_cfg.pixel_dim
+        value_dim = model_cfg.value_dim
+        sensory_dim = model_cfg.sensory_dim
+        final_dim = model_cfg.mask_encoder.final_dim
+
+        self.single_object = single_object
+        extra_dim = 1 if single_object else 2
+
+        # if model_cfg.pretrained_resnet is set in the model_cfg we get the value
+        # else default to True
+        is_pretrained_resnet = getattr(model_cfg,"pretrained_resnet",True)
+        if model_cfg.mask_encoder.type == 'resnet18':
+            network = resnet.resnet18(pretrained=is_pretrained_resnet, extra_dim=extra_dim)
+        elif model_cfg.mask_encoder.type == 'resnet50':
+            network = resnet.resnet50(pretrained=is_pretrained_resnet, extra_dim=extra_dim)
+        else:
+            raise NotImplementedError
+        self.conv1 = network.conv1
+        self.bn1 = network.bn1
+        self.relu = network.relu
+        self.maxpool = network.maxpool
+
+        self.layer1 = network.layer1
+        self.layer2 = network.layer2
+        self.layer3 = network.layer3
+
+        self.distributor = MainToGroupDistributor()
+        self.fuser = GroupFeatureFusionBlock(pixel_dim, final_dim, value_dim)
+
+        self.sensory_update = SensoryDeepUpdater(value_dim, sensory_dim)
+
+    def forward(self,
+                image: torch.Tensor,
+                pix_feat: torch.Tensor,
+                sensory: torch.Tensor,
+                masks: torch.Tensor,
+                others: torch.Tensor,
+                *,
+                deep_update: bool = True,
+                chunk_size: int = -1) -> (torch.Tensor, torch.Tensor):
+        # ms_features are from the key encoder
+        # we only use the first one (lowest resolution), following XMem
+        if self.single_object:
+            g = masks.unsqueeze(2)
+        else:
+            g = torch.stack([masks, others], dim=2)
+
+        g = self.distributor(image, g)
+
+        batch_size, num_objects = g.shape[:2]
+        if chunk_size < 1 or chunk_size >= num_objects:
+            chunk_size = num_objects
+            fast_path = True
+            new_sensory = sensory
+        else:
+            if deep_update:
+                new_sensory = torch.empty_like(sensory)
+            else:
+                new_sensory = sensory
+            fast_path = False
+
+        # chunk-by-chunk inference
+        all_g = []
+        for i in range(0, num_objects, chunk_size):
+            if fast_path:
+                g_chunk = g
+            else:
+                g_chunk = g[:, i:i + chunk_size]
+            actual_chunk_size = g_chunk.shape[1]
+            g_chunk = g_chunk.flatten(start_dim=0, end_dim=1)
+
+            g_chunk = self.conv1(g_chunk)
+            g_chunk = self.bn1(g_chunk)      # 1/2, 64
+            g_chunk = self.maxpool(g_chunk)  # 1/4, 64
+            g_chunk = self.relu(g_chunk)
+
+            g_chunk = self.layer1(g_chunk)   # 1/4
+            g_chunk = self.layer2(g_chunk)   # 1/8
+            g_chunk = self.layer3(g_chunk)   # 1/16
+
+            g_chunk = g_chunk.view(batch_size, actual_chunk_size, *g_chunk.shape[1:])
+            g_chunk = self.fuser(pix_feat, g_chunk)
+            all_g.append(g_chunk)
+            if deep_update:
+                if fast_path:
+                    new_sensory = self.sensory_update(g_chunk, sensory)
+                else:
+                    new_sensory[:, i:i + chunk_size] = self.sensory_update(
+                        g_chunk, sensory[:, i:i + chunk_size])
+        g = torch.cat(all_g, dim=1)
+
+        return g, new_sensory
+
+    # override the default train() to freeze BN statistics
+    def train(self, mode=True):
+        self.training = False
+        for module in self.children():
+            module.train(False)
+        return self
+
+
+class PixelFeatureFuser(nn.Module):
+    def __init__(self, model_cfg: DictConfig, single_object=False):
+        super().__init__()
+        value_dim = model_cfg.value_dim
+        sensory_dim = model_cfg.sensory_dim
+        pixel_dim = model_cfg.pixel_dim
+        embed_dim = model_cfg.embed_dim
+        self.single_object = single_object
+
+        self.fuser = GroupFeatureFusionBlock(pixel_dim, value_dim, embed_dim)
+        if self.single_object:
+            self.sensory_compress = GConv2d(sensory_dim + 1, value_dim, kernel_size=1)
+        else:
+            self.sensory_compress = GConv2d(sensory_dim + 2, value_dim, kernel_size=1)
+
+    def forward(self,
+                pix_feat: torch.Tensor,
+                pixel_memory: torch.Tensor,
+                sensory_memory: torch.Tensor,
+                last_mask: torch.Tensor,
+                last_others: torch.Tensor,
+                *,
+                chunk_size: int = -1) -> torch.Tensor:
+        batch_size, num_objects = pixel_memory.shape[:2]
+
+        if self.single_object:
+            last_mask = last_mask.unsqueeze(2)
+        else:
+            last_mask = torch.stack([last_mask, last_others], dim=2)
+
+        if chunk_size < 1:
+            chunk_size = num_objects
+
+        # chunk-by-chunk inference
+        all_p16 = []
+        for i in range(0, num_objects, chunk_size):
+            sensory_readout = self.sensory_compress(
+                torch.cat([sensory_memory[:, i:i + chunk_size], last_mask[:, i:i + chunk_size]], 2))
+            p16 = pixel_memory[:, i:i + chunk_size] + sensory_readout
+            p16 = self.fuser(pix_feat, p16)
+            all_p16.append(p16)
+        p16 = torch.cat(all_p16, dim=1)
+
+        return p16
+
+
+class MaskDecoder(nn.Module):
+    def __init__(self, model_cfg: DictConfig):
+        super().__init__()
+        embed_dim = model_cfg.embed_dim
+        sensory_dim = model_cfg.sensory_dim
+        ms_image_dims = model_cfg.pixel_encoder.ms_dims
+        up_dims = model_cfg.mask_decoder.up_dims
+
+        assert embed_dim == up_dims[0]
+
+        self.sensory_update = SensoryUpdater_fullscale([up_dims[0], up_dims[1], up_dims[2], up_dims[3], up_dims[4] + 1], sensory_dim,
+                                             sensory_dim)
+
+        self.decoder_feat_proc = DecoderFeatureProcessor(ms_image_dims[1:], up_dims[:-1])
+        self.up_16_8 = MaskUpsampleBlock(up_dims[0], up_dims[1])
+        self.up_8_4 = MaskUpsampleBlock(up_dims[1], up_dims[2])
+        # newly add for alpha matte
+        self.up_4_2 = MaskUpsampleBlock(up_dims[2], up_dims[3])
+        self.up_2_1 = MaskUpsampleBlock(up_dims[3], up_dims[4])
+
+        self.pred_seg = nn.Conv2d(up_dims[-1], 1, kernel_size=3, padding=1)
+        self.pred_mat = nn.Conv2d(up_dims[-1], 1, kernel_size=3, padding=1)
+
+    def forward(self,
+                ms_image_feat: Iterable[torch.Tensor],
+                memory_readout: torch.Tensor,
+                sensory: torch.Tensor,
+                *,
+                chunk_size: int = -1,
+                update_sensory: bool = True,
+                seg_pass: bool = False,
+                last_mask=None,
+                sigmoid_residual=False) -> (torch.Tensor, torch.Tensor):
+
+        batch_size, num_objects = memory_readout.shape[:2]
+        f8, f4, f2, f1 = self.decoder_feat_proc(ms_image_feat[1:])
+        if chunk_size < 1 or chunk_size >= num_objects:
+            chunk_size = num_objects
+            fast_path = True
+            new_sensory = sensory
+        else:
+            if update_sensory:
+                new_sensory = torch.empty_like(sensory)
+            else:
+                new_sensory = sensory
+            fast_path = False
+
+        # chunk-by-chunk inference
+        all_logits = []
+        for i in range(0, num_objects, chunk_size):
+            if fast_path:
+                p16 = memory_readout
+            else:
+                p16 = memory_readout[:, i:i + chunk_size]
+            actual_chunk_size = p16.shape[1]
+
+            p8 = self.up_16_8(p16, f8)
+            p4 = self.up_8_4(p8, f4)
+            p2 = self.up_4_2(p4, f2)
+            p1 = self.up_2_1(p2, f1)
+            with torch.amp.autocast("cuda",enabled=False):
+                if seg_pass:
+                    if last_mask is not None:
+                        res = self.pred_seg(F.relu(p1.flatten(start_dim=0, end_dim=1).float()))
+                        if sigmoid_residual:
+                            res = (torch.sigmoid(res) - 0.5) * 2  # regularization: (-1, 1) change on last mask
+                        logits = last_mask + res
+                    else:
+                        logits = self.pred_seg(F.relu(p1.flatten(start_dim=0, end_dim=1).float()))
+                else:
+                    if last_mask is not None:
+                        res = self.pred_mat(F.relu(p1.flatten(start_dim=0, end_dim=1).float()))
+                        if sigmoid_residual:
+                            res = (torch.sigmoid(res) - 0.5) * 2  # regularization: (-1, 1) change on last mask
+                        logits = last_mask + res
+                    else:
+                        logits = self.pred_mat(F.relu(p1.flatten(start_dim=0, end_dim=1).float()))
+            ## SensoryUpdater_fullscale
+            if update_sensory:
+                p1 = torch.cat(
+                    [p1, logits.view(batch_size, actual_chunk_size, 1, *logits.shape[-2:])], 2)
+                if fast_path:
+                    new_sensory = self.sensory_update([p16, p8, p4, p2, p1], sensory)
+                else:
+                    new_sensory[:,
+                                i:i + chunk_size] = self.sensory_update([p16, p8, p4, p2, p1],
+                                                                        sensory[:,
+                                                                                i:i + chunk_size])
+            all_logits.append(logits)
+        logits = torch.cat(all_logits, dim=0)
+        logits = logits.view(batch_size, num_objects, *logits.shape[-2:])
+
+        return new_sensory, logits

utility/matanyone2/matanyone/model/channel_attn.py

@@ -0,0 +1,39 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class CAResBlock(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int, residual: bool = True):
+        super().__init__()
+        self.residual = residual
+        self.conv1 = nn.Conv2d(in_dim, out_dim, kernel_size=3, padding=1)
+        self.conv2 = nn.Conv2d(out_dim, out_dim, kernel_size=3, padding=1)
+
+        t = int((abs(math.log2(out_dim)) + 1) // 2)
+        k = t if t % 2 else t + 1
+        self.pool = nn.AdaptiveAvgPool2d(1)
+        self.conv = nn.Conv1d(1, 1, kernel_size=k, padding=(k - 1) // 2, bias=False)
+
+        if self.residual:
+            if in_dim == out_dim:
+                self.downsample = nn.Identity()
+            else:
+                self.downsample = nn.Conv2d(in_dim, out_dim, kernel_size=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r = x
+        x = self.conv1(F.relu(x))
+        x = self.conv2(F.relu(x))
+
+        b, c = x.shape[:2]
+        w = self.pool(x).view(b, 1, c)
+        w = self.conv(w).transpose(-1, -2).unsqueeze(-1).sigmoid()  # B*C*1*1
+
+        if self.residual:
+            x = x * w + self.downsample(r)
+        else:
+            x = x * w
+
+        return x

utility/matanyone2/matanyone/model/group_modules.py

@@ -0,0 +1,127 @@
+from typing import Optional
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from matanyone.model.channel_attn import CAResBlock
+
+def interpolate_groups(g: torch.Tensor, ratio: float, mode: str,
+                       align_corners: bool) -> torch.Tensor:
+    batch_size, num_objects = g.shape[:2]
+    g = F.interpolate(g.flatten(start_dim=0, end_dim=1),
+                      scale_factor=ratio,
+                      mode=mode,
+                      align_corners=align_corners)
+    g = g.view(batch_size, num_objects, *g.shape[1:])
+    return g
+
+
+def upsample_groups(g: torch.Tensor,
+                    ratio: float = 2,
+                    mode: str = 'bilinear',
+                    align_corners: bool = False) -> torch.Tensor:
+    return interpolate_groups(g, ratio, mode, align_corners)
+
+
+def downsample_groups(g: torch.Tensor,
+                      ratio: float = 1 / 2,
+                      mode: str = 'area',
+                      align_corners: bool = None) -> torch.Tensor:
+    return interpolate_groups(g, ratio, mode, align_corners)
+
+
+class GConv2d(nn.Conv2d):
+    def forward(self, g: torch.Tensor) -> torch.Tensor:
+        batch_size, num_objects = g.shape[:2]
+        g = super().forward(g.flatten(start_dim=0, end_dim=1))
+        return g.view(batch_size, num_objects, *g.shape[1:])
+
+
+class GroupResBlock(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int):
+        super().__init__()
+
+        if in_dim == out_dim:
+            self.downsample = nn.Identity()
+        else:
+            self.downsample = GConv2d(in_dim, out_dim, kernel_size=1)
+
+        self.conv1 = GConv2d(in_dim, out_dim, kernel_size=3, padding=1)
+        self.conv2 = GConv2d(out_dim, out_dim, kernel_size=3, padding=1)
+
+    def forward(self, g: torch.Tensor) -> torch.Tensor:
+        out_g = self.conv1(F.relu(g))
+        out_g = self.conv2(F.relu(out_g))
+
+        g = self.downsample(g)
+
+        return out_g + g
+
+
+class MainToGroupDistributor(nn.Module):
+    def __init__(self,
+                 x_transform: Optional[nn.Module] = None,
+                 g_transform: Optional[nn.Module] = None,
+                 method: str = 'cat',
+                 reverse_order: bool = False):
+        super().__init__()
+
+        self.x_transform = x_transform
+        self.g_transform = g_transform
+        self.method = method
+        self.reverse_order = reverse_order
+
+    def forward(self, x: torch.Tensor, g: torch.Tensor, skip_expand: bool = False) -> torch.Tensor:
+        num_objects = g.shape[1]
+
+        if self.x_transform is not None:
+            x = self.x_transform(x)
+
+        if self.g_transform is not None:
+            g = self.g_transform(g)
+
+        if not skip_expand:
+            x = x.unsqueeze(1).expand(-1, num_objects, -1, -1, -1)
+
+        if self.method == 'cat':
+            if self.reverse_order:
+                g = torch.cat([g, x], 2)
+            else:
+                g = torch.cat([x, g], 2)
+        elif self.method == 'add':
+            g = x + g
+        elif self.method == 'mulcat':
+            g = torch.cat([x * g, g], dim=2)
+        elif self.method == 'muladd':
+            g = x * g + g
+        else:
+            raise NotImplementedError
+
+        return g
+
+
+class GroupFeatureFusionBlock(nn.Module):
+    def __init__(self, x_in_dim: int, g_in_dim: int, out_dim: int):
+        super().__init__()
+
+        x_transform = nn.Conv2d(x_in_dim, out_dim, kernel_size=1)
+        g_transform = GConv2d(g_in_dim, out_dim, kernel_size=1)
+
+        self.distributor = MainToGroupDistributor(x_transform=x_transform,
+                                                  g_transform=g_transform,
+                                                  method='add')
+        self.block1 = CAResBlock(out_dim, out_dim)
+        self.block2 = CAResBlock(out_dim, out_dim)
+
+    def forward(self, x: torch.Tensor, g: torch.Tensor) -> torch.Tensor:
+        batch_size, num_objects = g.shape[:2]
+
+        g = self.distributor(x, g)
+
+        g = g.flatten(start_dim=0, end_dim=1)
+
+        g = self.block1(g)
+        g = self.block2(g)
+
+        g = g.view(batch_size, num_objects, *g.shape[1:])
+
+        return g

utility/matanyone2/matanyone/model/matanyone.py

@@ -0,0 +1,332 @@
+from typing import List, Dict, Iterable, Tuple
+import logging
+from omegaconf import DictConfig
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from omegaconf import OmegaConf
+from huggingface_hub import PyTorchModelHubMixin
+
+from matanyone.model.big_modules import PixelEncoder, UncertPred, KeyProjection, MaskEncoder, PixelFeatureFuser, MaskDecoder
+from matanyone.model.aux_modules import AuxComputer
+from matanyone.model.utils.memory_utils import get_affinity, readout
+from matanyone.model.transformer.object_transformer import QueryTransformer
+from matanyone.model.transformer.object_summarizer import ObjectSummarizer
+from matanyone.utils.tensor_utils import aggregate
+
+log = logging.getLogger()
+class MatAnyone(nn.Module,
+                PyTorchModelHubMixin,
+                library_name="matanyone",
+                repo_url="https://github.com/pq-yang/MatAnyone",
+                coders={
+                    DictConfig: (
+                        lambda x: OmegaConf.to_container(x),
+                        lambda data: OmegaConf.create(data),
+                    )
+                },
+        ):
+
+    def __init__(self, cfg: DictConfig, *, single_object=False):
+        super().__init__()
+        self.cfg = cfg
+        model_cfg = cfg.model
+        self.ms_dims = model_cfg.pixel_encoder.ms_dims
+        self.key_dim = model_cfg.key_dim
+        self.value_dim = model_cfg.value_dim
+        self.sensory_dim = model_cfg.sensory_dim
+        self.pixel_dim = model_cfg.pixel_dim
+        self.embed_dim = model_cfg.embed_dim
+        self.single_object = single_object
+
+        self.pixel_encoder = PixelEncoder(model_cfg)
+        self.pix_feat_proj = nn.Conv2d(self.ms_dims[0], self.pixel_dim, kernel_size=1)
+        self.key_proj = KeyProjection(model_cfg)
+        self.mask_encoder = MaskEncoder(model_cfg, single_object=single_object)
+        self.mask_decoder = MaskDecoder(model_cfg)
+        self.pixel_fuser = PixelFeatureFuser(model_cfg, single_object=single_object)
+        self.object_transformer = QueryTransformer(model_cfg)
+        self.object_summarizer = ObjectSummarizer(model_cfg)
+        self.aux_computer = AuxComputer(cfg)
+        self.temp_sparity = UncertPred(model_cfg)
+
+        self.register_buffer("pixel_mean", torch.Tensor(model_cfg.pixel_mean).view(-1, 1, 1), False)
+        self.register_buffer("pixel_std", torch.Tensor(model_cfg.pixel_std).view(-1, 1, 1), False)
+
+    def _get_others(self, masks: torch.Tensor) -> torch.Tensor:
+        # for each object, return the sum of masks of all other objects
+        if self.single_object:
+            return None
+
+        num_objects = masks.shape[1]
+        if num_objects >= 1:
+            others = (masks.sum(dim=1, keepdim=True) - masks).clamp(0, 1)
+        else:
+            others = torch.zeros_like(masks)
+        return others
+
+    def pred_uncertainty(self, last_pix_feat: torch.Tensor, cur_pix_feat: torch.Tensor, last_mask: torch.Tensor, mem_val_diff:torch.Tensor):
+        logits = self.temp_sparity(last_frame_feat=last_pix_feat,
+                                cur_frame_feat=cur_pix_feat,
+                                last_mask=last_mask,
+                                mem_val_diff=mem_val_diff)
+
+        prob = torch.sigmoid(logits)
+        mask = (prob > 0) + 0
+
+        uncert_output = {"logits": logits,
+                     "prob": prob,
+                     "mask": mask}
+
+        return uncert_output
+
+    def encode_image(self, image: torch.Tensor, seq_length=None, last_feats=None) -> (Iterable[torch.Tensor], torch.Tensor): # type: ignore
+        image = (image - self.pixel_mean) / self.pixel_std
+        ms_image_feat = self.pixel_encoder(image, seq_length) # f16, f8, f4, f2, f1
+        return ms_image_feat, self.pix_feat_proj(ms_image_feat[0])
+
+    def encode_mask(
+            self,
+            image: torch.Tensor,
+            ms_features: List[torch.Tensor],
+            sensory: torch.Tensor,
+            masks: torch.Tensor,
+            *,
+            deep_update: bool = True,
+            chunk_size: int = -1,
+            need_weights: bool = False) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        image = (image - self.pixel_mean) / self.pixel_std
+        others = self._get_others(masks)
+
+        mask_value, new_sensory = self.mask_encoder(image,
+                                                    ms_features,
+                                                    sensory,
+                                                    masks,
+                                                    others,
+                                                    deep_update=deep_update,
+                                                    chunk_size=chunk_size)
+        object_summaries, object_logits = self.object_summarizer(masks, mask_value, need_weights)
+        return mask_value, new_sensory, object_summaries, object_logits
+
+    def transform_key(self,
+                      final_pix_feat: torch.Tensor,
+                      *,
+                      need_sk: bool = True,
+                      need_ek: bool = True) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        key, shrinkage, selection = self.key_proj(final_pix_feat, need_s=need_sk, need_e=need_ek)
+        return key, shrinkage, selection
+
+    # Used in training only.
+    # This step is replaced by MemoryManager in test time
+    def read_memory(self, query_key: torch.Tensor, query_selection: torch.Tensor,
+                    memory_key: torch.Tensor, memory_shrinkage: torch.Tensor,
+                    msk_value: torch.Tensor, obj_memory: torch.Tensor, pix_feat: torch.Tensor,
+                    sensory: torch.Tensor, last_mask: torch.Tensor,
+                    selector: torch.Tensor, uncert_output=None, seg_pass=False,
+                    last_pix_feat=None, last_pred_mask=None) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+        """
+        query_key       : B * CK * H * W
+        query_selection : B * CK * H * W
+        memory_key      : B * CK * T * H * W
+        memory_shrinkage: B * 1  * T * H * W
+        msk_value       : B * num_objects * CV * T * H * W
+        obj_memory      : B * num_objects * T * num_summaries * C
+        pixel_feature   : B * C * H * W
+        """
+        batch_size, num_objects = msk_value.shape[:2]
+
+        uncert_mask = uncert_output["mask"] if uncert_output is not None else None
+
+        # read using visual attention
+        with torch.amp.autocast("cuda",enabled=False):
+            affinity = get_affinity(memory_key.float(), memory_shrinkage.float(), query_key.float(),
+                                    query_selection.float(), uncert_mask=uncert_mask)
+
+            msk_value = msk_value.flatten(start_dim=1, end_dim=2).float()
+
+            # B * (num_objects*CV) * H * W
+            pixel_readout = readout(affinity, msk_value, uncert_mask)
+            pixel_readout = pixel_readout.view(batch_size, num_objects, self.value_dim,
+                                            *pixel_readout.shape[-2:])
+            
+            uncert_output = self.pred_uncertainty(last_pix_feat, pix_feat, last_pred_mask, pixel_readout[:,0]-msk_value[:,:,-1])
+            uncert_prob = uncert_output["prob"].unsqueeze(1) # b n 1 h w
+            pixel_readout = pixel_readout*uncert_prob + msk_value[:,:,-1].unsqueeze(1)*(1-uncert_prob)
+                
+        pixel_readout = self.pixel_fusion(pix_feat, pixel_readout, sensory, last_mask)
+
+
+        # read from query transformer
+        mem_readout, aux_features = self.readout_query(pixel_readout, obj_memory, selector=selector, seg_pass=seg_pass)
+
+        aux_output = {
+            'sensory': sensory,
+            'q_logits': aux_features['logits'] if aux_features else None,
+            'attn_mask': aux_features['attn_mask'] if aux_features else None,
+        }
+
+        return mem_readout, aux_output, uncert_output
+    
+    def read_first_frame_memory(self, pixel_readout,
+                    obj_memory: torch.Tensor, pix_feat: torch.Tensor,
+                    sensory: torch.Tensor, last_mask: torch.Tensor,
+                    selector: torch.Tensor, seg_pass=False) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+        """
+        query_key       : B * CK * H * W
+        query_selection : B * CK * H * W
+        memory_key      : B * CK * T * H * W
+        memory_shrinkage: B * 1  * T * H * W
+        msk_value       : B * num_objects * CV * T * H * W
+        obj_memory      : B * num_objects * T * num_summaries * C
+        pixel_feature   : B * C * H * W
+        """
+                
+        pixel_readout = self.pixel_fusion(pix_feat, pixel_readout, sensory, last_mask)
+
+        # read from query transformer
+        mem_readout, aux_features = self.readout_query(pixel_readout, obj_memory, selector=selector, seg_pass=seg_pass)
+
+        aux_output = {
+            'sensory': sensory,
+            'q_logits': aux_features['logits'] if aux_features else None,
+            'attn_mask': aux_features['attn_mask'] if aux_features else None,
+        }
+
+        return mem_readout, aux_output
+
+    def pixel_fusion(self,
+                     pix_feat: torch.Tensor,
+                     pixel: torch.Tensor,
+                     sensory: torch.Tensor,
+                     last_mask: torch.Tensor,
+                     *,
+                     chunk_size: int = -1) -> torch.Tensor:
+        last_mask = F.interpolate(last_mask, size=sensory.shape[-2:], mode='area')
+        last_others = self._get_others(last_mask)
+        fused = self.pixel_fuser(pix_feat,
+                                 pixel,
+                                 sensory,
+                                 last_mask,
+                                 last_others,
+                                 chunk_size=chunk_size)
+        return fused
+
+    def readout_query(self,
+                      pixel_readout,
+                      obj_memory,
+                      *,
+                      selector=None,
+                      need_weights=False,
+                      seg_pass=False) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+        return self.object_transformer(pixel_readout,
+                                       obj_memory,
+                                       selector=selector,
+                                       need_weights=need_weights,
+                                       seg_pass=seg_pass)
+
+    def segment(self,
+                ms_image_feat: List[torch.Tensor],
+                memory_readout: torch.Tensor,
+                sensory: torch.Tensor,
+                *,
+                selector: bool = None,
+                chunk_size: int = -1,
+                update_sensory: bool = True,
+                seg_pass: bool = False,
+                clamp_mat: bool = True,
+                last_mask=None,
+                sigmoid_residual=False,
+                seg_mat=False) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        multi_scale_features is from the key encoder for skip-connection
+        memory_readout is from working/long-term memory
+        sensory is the sensory memory
+        last_mask is the mask from the last frame, supplementing sensory memory
+        selector is 1 if an object exists, and 0 otherwise. We use it to filter padded objects
+            during training.
+        """
+        #### use mat head for seg data
+        if seg_mat:
+            assert seg_pass
+            seg_pass = False
+        ####
+        sensory, logits = self.mask_decoder(ms_image_feat,
+                                        memory_readout,
+                                        sensory,
+                                        chunk_size=chunk_size,
+                                        update_sensory=update_sensory,
+                                        seg_pass = seg_pass,
+                                        last_mask=last_mask,
+                                        sigmoid_residual=sigmoid_residual)
+        if seg_pass:
+            prob = torch.sigmoid(logits)
+            if selector is not None:
+                prob = prob * selector
+
+            # Softmax over all objects[]
+            logits = aggregate(prob, dim=1)
+            prob = F.softmax(logits, dim=1)
+        else:
+            if clamp_mat:
+                logits = logits.clamp(0.0, 1.0)
+            logits = torch.cat([torch.prod(1 - logits, dim=1, keepdim=True), logits], 1)
+            prob = logits
+        
+        return sensory, logits, prob
+
+    def compute_aux(self, pix_feat: torch.Tensor, aux_inputs: Dict[str, torch.Tensor],
+                    selector: torch.Tensor, seg_pass=False) -> Dict[str, torch.Tensor]:
+        return self.aux_computer(pix_feat, aux_inputs, selector, seg_pass=seg_pass)
+
+    def forward(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def load_weights(self, src_dict, init_as_zero_if_needed=False) -> None:
+        if not self.single_object:
+            # Map single-object weight to multi-object weight (4->5 out channels in conv1)
+            for k in list(src_dict.keys()):
+                if k == 'mask_encoder.conv1.weight':
+                    if src_dict[k].shape[1] == 4:
+                        log.info(f'Converting {k} from single object to multiple objects.')
+                        pads = torch.zeros((64, 1, 7, 7), device=src_dict[k].device)
+                        if not init_as_zero_if_needed:
+                            nn.init.orthogonal_(pads)
+                            log.info(f'Randomly initialized padding for {k}.')
+                        else:
+                            log.info(f'Zero-initialized padding for {k}.')
+                        src_dict[k] = torch.cat([src_dict[k], pads], 1)
+                elif k == 'pixel_fuser.sensory_compress.weight':
+                    if src_dict[k].shape[1] == self.sensory_dim + 1:
+                        log.info(f'Converting {k} from single object to multiple objects.')
+                        pads = torch.zeros((self.value_dim, 1, 1, 1), device=src_dict[k].device)
+                        if not init_as_zero_if_needed:
+                            nn.init.orthogonal_(pads)
+                            log.info(f'Randomly initialized padding for {k}.')
+                        else:
+                            log.info(f'Zero-initialized padding for {k}.')
+                        src_dict[k] = torch.cat([src_dict[k], pads], 1)
+        elif self.single_object:
+            """
+            If the model is multiple-object and we are training in single-object, 
+            we strip the last channel of conv1.
+            This is not supposed to happen in standard training except when users are trying to
+            finetune a trained model with single object datasets.
+            """
+            if src_dict['mask_encoder.conv1.weight'].shape[1] == 5:
+                log.warning('Converting mask_encoder.conv1.weight from multiple objects to single object.'
+                            'This is not supposed to happen in standard training.')
+                src_dict['mask_encoder.conv1.weight'] = src_dict['mask_encoder.conv1.weight'][:, :-1]
+                src_dict['pixel_fuser.sensory_compress.weight'] = src_dict['pixel_fuser.sensory_compress.weight'][:, :-1]
+
+        for k in src_dict:
+            if k not in self.state_dict():
+                log.info(f'Key {k} found in src_dict but not in self.state_dict()!!!')
+        for k in self.state_dict():
+            if k not in src_dict:
+                log.info(f'Key {k} found in self.state_dict() but not in src_dict!!!')
+                
+        self.load_state_dict(src_dict, strict=False)
+
+    @property
+    def device(self) -> torch.device:
+        return self.pixel_mean.device

utility/matanyone2/matanyone/model/modules.py

@@ -0,0 +1,149 @@
+from typing import List, Iterable
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from matanyone.model.group_modules import MainToGroupDistributor, GroupResBlock, upsample_groups, GConv2d, downsample_groups
+
+
+class UpsampleBlock(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int, scale_factor: int = 2):
+        super().__init__()
+        self.out_conv = ResBlock(in_dim, out_dim)
+        self.scale_factor = scale_factor
+
+    def forward(self, in_g: torch.Tensor, skip_f: torch.Tensor) -> torch.Tensor:
+        g = F.interpolate(in_g,
+                      scale_factor=self.scale_factor,
+                      mode='bilinear')
+        g = self.out_conv(g)
+        g = g + skip_f
+        return g
+
+class MaskUpsampleBlock(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int, scale_factor: int = 2):
+        super().__init__()
+        self.distributor = MainToGroupDistributor(method='add')
+        self.out_conv = GroupResBlock(in_dim, out_dim)
+        self.scale_factor = scale_factor
+
+    def forward(self, in_g: torch.Tensor, skip_f: torch.Tensor) -> torch.Tensor:
+        g = upsample_groups(in_g, ratio=self.scale_factor)
+        g = self.distributor(skip_f, g)
+        g = self.out_conv(g)
+        return g
+    
+
+class DecoderFeatureProcessor(nn.Module):
+    def __init__(self, decoder_dims: List[int], out_dims: List[int]):
+        super().__init__()
+        self.transforms = nn.ModuleList([
+            nn.Conv2d(d_dim, p_dim, kernel_size=1) for d_dim, p_dim in zip(decoder_dims, out_dims)
+        ])
+
+    def forward(self, multi_scale_features: Iterable[torch.Tensor]) -> List[torch.Tensor]:
+        outputs = [func(x) for x, func in zip(multi_scale_features, self.transforms)]
+        return outputs
+
+
+# @torch.jit.script
+def _recurrent_update(h: torch.Tensor, values: torch.Tensor) -> torch.Tensor:
+    # h: batch_size * num_objects * hidden_dim * h * w
+    # values: batch_size * num_objects * (hidden_dim*3) * h * w
+    dim = values.shape[2] // 3
+    forget_gate = torch.sigmoid(values[:, :, :dim])
+    update_gate = torch.sigmoid(values[:, :, dim:dim * 2])
+    new_value = torch.tanh(values[:, :, dim * 2:])
+    new_h = forget_gate * h * (1 - update_gate) + update_gate * new_value
+    return new_h
+
+
+class SensoryUpdater_fullscale(nn.Module):
+    # Used in the decoder, multi-scale feature + GRU
+    def __init__(self, g_dims: List[int], mid_dim: int, sensory_dim: int):
+        super().__init__()
+        self.g16_conv = GConv2d(g_dims[0], mid_dim, kernel_size=1)
+        self.g8_conv = GConv2d(g_dims[1], mid_dim, kernel_size=1)
+        self.g4_conv = GConv2d(g_dims[2], mid_dim, kernel_size=1)
+        self.g2_conv = GConv2d(g_dims[3], mid_dim, kernel_size=1)
+        self.g1_conv = GConv2d(g_dims[4], mid_dim, kernel_size=1)
+
+        self.transform = GConv2d(mid_dim + sensory_dim, sensory_dim * 3, kernel_size=3, padding=1)
+
+        nn.init.xavier_normal_(self.transform.weight)
+
+    def forward(self, g: torch.Tensor, h: torch.Tensor) -> torch.Tensor:
+        g = self.g16_conv(g[0]) + self.g8_conv(downsample_groups(g[1], ratio=1/2)) + \
+            self.g4_conv(downsample_groups(g[2], ratio=1/4)) + \
+            self.g2_conv(downsample_groups(g[3], ratio=1/8)) + \
+            self.g1_conv(downsample_groups(g[4], ratio=1/16))
+
+        with torch.amp.autocast("cuda",enabled=False):
+            g = g.float()
+            h = h.float()
+            values = self.transform(torch.cat([g, h], dim=2))
+            new_h = _recurrent_update(h, values)
+
+        return new_h
+
+class SensoryUpdater(nn.Module):
+    # Used in the decoder, multi-scale feature + GRU
+    def __init__(self, g_dims: List[int], mid_dim: int, sensory_dim: int):
+        super().__init__()
+        self.g16_conv = GConv2d(g_dims[0], mid_dim, kernel_size=1)
+        self.g8_conv = GConv2d(g_dims[1], mid_dim, kernel_size=1)
+        self.g4_conv = GConv2d(g_dims[2], mid_dim, kernel_size=1)
+
+        self.transform = GConv2d(mid_dim + sensory_dim, sensory_dim * 3, kernel_size=3, padding=1)
+
+        nn.init.xavier_normal_(self.transform.weight)
+
+    def forward(self, g: torch.Tensor, h: torch.Tensor) -> torch.Tensor:
+        g = self.g16_conv(g[0]) + self.g8_conv(downsample_groups(g[1], ratio=1/2)) + \
+            self.g4_conv(downsample_groups(g[2], ratio=1/4))
+
+        with torch.amp.autocast("cuda",enabled=False):
+            g = g.float()
+            h = h.float()
+            values = self.transform(torch.cat([g, h], dim=2))
+            new_h = _recurrent_update(h, values)
+
+        return new_h
+
+
+class SensoryDeepUpdater(nn.Module):
+    def __init__(self, f_dim: int, sensory_dim: int):
+        super().__init__()
+        self.transform = GConv2d(f_dim + sensory_dim, sensory_dim * 3, kernel_size=3, padding=1)
+
+        nn.init.xavier_normal_(self.transform.weight)
+
+    def forward(self, g: torch.Tensor, h: torch.Tensor) -> torch.Tensor:
+        with torch.amp.autocast("cuda",enabled=False):
+            g = g.float()
+            h = h.float()
+            values = self.transform(torch.cat([g, h], dim=2))
+            new_h = _recurrent_update(h, values)
+
+        return new_h
+
+  
+class ResBlock(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int):
+        super().__init__()
+
+        if in_dim == out_dim:
+            self.downsample = nn.Identity()
+        else:
+            self.downsample = nn.Conv2d(in_dim, out_dim, kernel_size=1)
+
+        self.conv1 = nn.Conv2d(in_dim, out_dim, kernel_size=3, padding=1)
+        self.conv2 = nn.Conv2d(out_dim, out_dim, kernel_size=3, padding=1)
+
+    def forward(self, g: torch.Tensor) -> torch.Tensor:
+        out_g = self.conv1(F.relu(g))
+        out_g = self.conv2(F.relu(out_g))
+
+        g = self.downsample(g)
+
+        return out_g + g

utility/matanyone2/matanyone/model/transformer/object_summarizer.py

@@ -0,0 +1,89 @@
+from typing import Optional
+from omegaconf import DictConfig
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from matanyone.model.transformer.positional_encoding import PositionalEncoding
+
+
+# @torch.jit.script
+def _weighted_pooling(masks: torch.Tensor, value: torch.Tensor,
+                      logits: torch.Tensor) -> (torch.Tensor, torch.Tensor):
+    # value: B*num_objects*H*W*value_dim
+    # logits: B*num_objects*H*W*num_summaries
+    # masks: B*num_objects*H*W*num_summaries: 1 if allowed
+    weights = logits.sigmoid() * masks
+    # B*num_objects*num_summaries*value_dim
+    sums = torch.einsum('bkhwq,bkhwc->bkqc', weights, value)
+    # B*num_objects*H*W*num_summaries -> B*num_objects*num_summaries*1
+    area = weights.flatten(start_dim=2, end_dim=3).sum(2).unsqueeze(-1)
+
+    # B*num_objects*num_summaries*value_dim
+    return sums, area
+
+
+class ObjectSummarizer(nn.Module):
+    def __init__(self, model_cfg: DictConfig):
+        super().__init__()
+
+        this_cfg = model_cfg.object_summarizer
+        self.value_dim = model_cfg.value_dim
+        self.embed_dim = this_cfg.embed_dim
+        self.num_summaries = this_cfg.num_summaries
+        self.add_pe = this_cfg.add_pe
+        self.pixel_pe_scale = model_cfg.pixel_pe_scale
+        self.pixel_pe_temperature = model_cfg.pixel_pe_temperature
+
+        if self.add_pe:
+            self.pos_enc = PositionalEncoding(self.embed_dim,
+                                              scale=self.pixel_pe_scale,
+                                              temperature=self.pixel_pe_temperature)
+
+        self.input_proj = nn.Linear(self.value_dim, self.embed_dim)
+        self.feature_pred = nn.Sequential(
+            nn.Linear(self.embed_dim, self.embed_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(self.embed_dim, self.embed_dim),
+        )
+        self.weights_pred = nn.Sequential(
+            nn.Linear(self.embed_dim, self.embed_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(self.embed_dim, self.num_summaries),
+        )
+
+    def forward(self,
+                masks: torch.Tensor,
+                value: torch.Tensor,
+                need_weights: bool = False) -> (torch.Tensor, Optional[torch.Tensor]):
+        # masks: B*num_objects*(H0)*(W0)
+        # value: B*num_objects*value_dim*H*W
+        # -> B*num_objects*H*W*value_dim
+        h, w = value.shape[-2:]
+        masks = F.interpolate(masks, size=(h, w), mode='area')
+        masks = masks.unsqueeze(-1)
+        inv_masks = 1 - masks
+        repeated_masks = torch.cat([
+            masks.expand(-1, -1, -1, -1, self.num_summaries // 2),
+            inv_masks.expand(-1, -1, -1, -1, self.num_summaries // 2),
+        ],
+                                   dim=-1)
+
+        value = value.permute(0, 1, 3, 4, 2)
+        value = self.input_proj(value)
+        if self.add_pe:
+            pe = self.pos_enc(value)
+            value = value + pe
+
+        with torch.amp.autocast("cuda",enabled=False):
+            value = value.float()
+            feature = self.feature_pred(value)
+            logits = self.weights_pred(value)
+            sums, area = _weighted_pooling(repeated_masks, feature, logits)
+
+        summaries = torch.cat([sums, area], dim=-1)
+
+        if need_weights:
+            return summaries, logits
+        else:
+            return summaries, None

utility/matanyone2/matanyone/model/transformer/object_transformer.py

@@ -0,0 +1,206 @@
+from typing import Dict, Optional
+from omegaconf import DictConfig
+
+import torch
+import torch.nn as nn
+from matanyone.model.group_modules import GConv2d
+from matanyone.utils.tensor_utils import aggregate
+from matanyone.model.transformer.positional_encoding import PositionalEncoding
+from matanyone.model.transformer.transformer_layers import *
+
+
+class QueryTransformerBlock(nn.Module):
+    def __init__(self, model_cfg: DictConfig):
+        super().__init__()
+
+        this_cfg = model_cfg.object_transformer
+        self.embed_dim = this_cfg.embed_dim
+        self.num_heads = this_cfg.num_heads
+        self.num_queries = this_cfg.num_queries
+        self.ff_dim = this_cfg.ff_dim
+
+        self.read_from_pixel = CrossAttention(self.embed_dim,
+                                              self.num_heads,
+                                              add_pe_to_qkv=this_cfg.read_from_pixel.add_pe_to_qkv)
+        self.self_attn = SelfAttention(self.embed_dim,
+                                       self.num_heads,
+                                       add_pe_to_qkv=this_cfg.query_self_attention.add_pe_to_qkv)
+        self.ffn = FFN(self.embed_dim, self.ff_dim)
+        self.read_from_query = CrossAttention(self.embed_dim,
+                                              self.num_heads,
+                                              add_pe_to_qkv=this_cfg.read_from_query.add_pe_to_qkv,
+                                              norm=this_cfg.read_from_query.output_norm)
+        self.pixel_ffn = PixelFFN(self.embed_dim)
+
+    def forward(
+            self,
+            x: torch.Tensor,
+            pixel: torch.Tensor,
+            query_pe: torch.Tensor,
+            pixel_pe: torch.Tensor,
+            attn_mask: torch.Tensor,
+            need_weights: bool = False) -> (torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor):
+        # x: (bs*num_objects)*num_queries*embed_dim
+        # pixel: bs*num_objects*C*H*W
+        # query_pe: (bs*num_objects)*num_queries*embed_dim
+        # pixel_pe: (bs*num_objects)*(H*W)*C
+        # attn_mask: (bs*num_objects*num_heads)*num_queries*(H*W)
+
+        # bs*num_objects*C*H*W -> (bs*num_objects)*(H*W)*C
+        pixel_flat = pixel.flatten(3, 4).flatten(0, 1).transpose(1, 2).contiguous()
+        x, q_weights = self.read_from_pixel(x,
+                                            pixel_flat,
+                                            query_pe,
+                                            pixel_pe,
+                                            attn_mask=attn_mask,
+                                            need_weights=need_weights)
+        x = self.self_attn(x, query_pe)
+        x = self.ffn(x)
+
+        pixel_flat, p_weights = self.read_from_query(pixel_flat,
+                                                     x,
+                                                     pixel_pe,
+                                                     query_pe,
+                                                     need_weights=need_weights)
+        pixel = self.pixel_ffn(pixel, pixel_flat)
+
+        if need_weights:
+            bs, num_objects, _, h, w = pixel.shape
+            q_weights = q_weights.view(bs, num_objects, self.num_heads, self.num_queries, h, w)
+            p_weights = p_weights.transpose(2, 3).view(bs, num_objects, self.num_heads,
+                                                       self.num_queries, h, w)
+
+        return x, pixel, q_weights, p_weights
+
+
+class QueryTransformer(nn.Module):
+    def __init__(self, model_cfg: DictConfig):
+        super().__init__()
+
+        this_cfg = model_cfg.object_transformer
+        self.value_dim = model_cfg.value_dim
+        self.embed_dim = this_cfg.embed_dim
+        self.num_heads = this_cfg.num_heads
+        self.num_queries = this_cfg.num_queries
+
+        # query initialization and embedding
+        self.query_init = nn.Embedding(self.num_queries, self.embed_dim)
+        self.query_emb = nn.Embedding(self.num_queries, self.embed_dim)
+
+        # projection from object summaries to query initialization and embedding
+        self.summary_to_query_init = nn.Linear(self.embed_dim, self.embed_dim)
+        self.summary_to_query_emb = nn.Linear(self.embed_dim, self.embed_dim)
+
+        self.pixel_pe_scale = model_cfg.pixel_pe_scale
+        self.pixel_pe_temperature = model_cfg.pixel_pe_temperature
+        self.pixel_init_proj = GConv2d(self.embed_dim, self.embed_dim, kernel_size=1)
+        self.pixel_emb_proj = GConv2d(self.embed_dim, self.embed_dim, kernel_size=1)
+        self.spatial_pe = PositionalEncoding(self.embed_dim,
+                                             scale=self.pixel_pe_scale,
+                                             temperature=self.pixel_pe_temperature,
+                                             channel_last=False,
+                                             transpose_output=True)
+
+        # transformer blocks
+        self.num_blocks = this_cfg.num_blocks
+        self.blocks = nn.ModuleList(
+            QueryTransformerBlock(model_cfg) for _ in range(self.num_blocks))
+        self.mask_pred = nn.ModuleList(
+            nn.Sequential(nn.ReLU(), GConv2d(self.embed_dim, 1, kernel_size=1))
+            for _ in range(self.num_blocks + 1))
+
+        self.act = nn.ReLU(inplace=True)
+
+    def forward(self,
+                pixel: torch.Tensor,
+                obj_summaries: torch.Tensor,
+                selector: Optional[torch.Tensor] = None,
+                need_weights: bool = False,
+                seg_pass=False) -> (torch.Tensor, Dict[str, torch.Tensor]):
+        # pixel: B*num_objects*embed_dim*H*W
+        # obj_summaries: B*num_objects*T*num_queries*embed_dim
+        T = obj_summaries.shape[2]
+        bs, num_objects, _, H, W = pixel.shape
+
+        # normalize object values
+        # the last channel is the cumulative area of the object
+        obj_summaries = obj_summaries.view(bs * num_objects, T, self.num_queries,
+                                           self.embed_dim + 1)
+        # sum over time
+        # during inference, T=1 as we already did streaming average in memory_manager
+        obj_sums = obj_summaries[:, :, :, :-1].sum(dim=1)
+        obj_area = obj_summaries[:, :, :, -1:].sum(dim=1)
+        obj_values = obj_sums / (obj_area + 1e-4)
+        obj_init = self.summary_to_query_init(obj_values)
+        obj_emb = self.summary_to_query_emb(obj_values)
+
+        # positional embeddings for object queries
+        query = self.query_init.weight.unsqueeze(0).expand(bs * num_objects, -1, -1) + obj_init
+        query_emb = self.query_emb.weight.unsqueeze(0).expand(bs * num_objects, -1, -1) + obj_emb
+
+        # positional embeddings for pixel features
+        pixel_init = self.pixel_init_proj(pixel)
+        pixel_emb = self.pixel_emb_proj(pixel)
+        pixel_pe = self.spatial_pe(pixel.flatten(0, 1))
+        pixel_emb = pixel_emb.flatten(3, 4).flatten(0, 1).transpose(1, 2).contiguous()
+        pixel_pe = pixel_pe.flatten(1, 2) + pixel_emb
+
+        pixel = pixel_init
+
+        # run the transformer
+        aux_features = {'logits': []}
+
+        # first aux output
+        aux_logits = self.mask_pred[0](pixel).squeeze(2)
+        attn_mask = self._get_aux_mask(aux_logits, selector, seg_pass=seg_pass)
+        aux_features['logits'].append(aux_logits)
+        for i in range(self.num_blocks):
+            query, pixel, q_weights, p_weights = self.blocks[i](query,
+                                                                pixel,
+                                                                query_emb,
+                                                                pixel_pe,
+                                                                attn_mask,
+                                                                need_weights=need_weights)
+
+            if self.training or i <= self.num_blocks - 1 or need_weights:
+                aux_logits = self.mask_pred[i + 1](pixel).squeeze(2)
+                attn_mask = self._get_aux_mask(aux_logits, selector, seg_pass=seg_pass)
+                aux_features['logits'].append(aux_logits)
+
+        aux_features['q_weights'] = q_weights  # last layer only
+        aux_features['p_weights'] = p_weights  # last layer only
+
+        if self.training:
+            # no need to save all heads
+            aux_features['attn_mask'] = attn_mask.view(bs, num_objects, self.num_heads,
+                                                       self.num_queries, H, W)[:, :, 0]
+
+        return pixel, aux_features
+
+    def _get_aux_mask(self, logits: torch.Tensor, selector: torch.Tensor, seg_pass=False) -> torch.Tensor:
+        # logits: batch_size*num_objects*H*W
+        # selector: batch_size*num_objects*1*1
+        # returns a mask of shape (batch_size*num_objects*num_heads)*num_queries*(H*W)
+        # where True means the attention is blocked
+
+        if selector is None:
+            prob = logits.sigmoid()
+        else:
+            prob = logits.sigmoid() * selector
+        logits = aggregate(prob, dim=1)
+
+        is_foreground = (logits[:, 1:] >= logits.max(dim=1, keepdim=True)[0])
+        foreground_mask = is_foreground.bool().flatten(start_dim=2)
+        inv_foreground_mask = ~foreground_mask
+        inv_background_mask = foreground_mask
+
+        aux_foreground_mask = inv_foreground_mask.unsqueeze(2).unsqueeze(2).repeat(
+            1, 1, self.num_heads, self.num_queries // 2, 1).flatten(start_dim=0, end_dim=2)
+        aux_background_mask = inv_background_mask.unsqueeze(2).unsqueeze(2).repeat(
+            1, 1, self.num_heads, self.num_queries // 2, 1).flatten(start_dim=0, end_dim=2)
+
+        aux_mask = torch.cat([aux_foreground_mask, aux_background_mask], dim=1)
+
+        aux_mask[torch.where(aux_mask.sum(-1) == aux_mask.shape[-1])] = False
+
+        return aux_mask

utility/matanyone2/matanyone/model/transformer/positional_encoding.py

@@ -0,0 +1,108 @@
+# Reference:
+# https://github.com/facebookresearch/Mask2Former/blob/main/mask2former/modeling/transformer_decoder/position_encoding.py
+# https://github.com/tatp22/multidim-positional-encoding/blob/master/positional_encodings/torch_encodings.py
+
+import math
+
+import numpy as np
+import torch
+from torch import nn
+
+
+def get_emb(sin_inp: torch.Tensor) -> torch.Tensor:
+    """
+    Gets a base embedding for one dimension with sin and cos intertwined
+    """
+    emb = torch.stack((sin_inp.sin(), sin_inp.cos()), dim=-1)
+    return torch.flatten(emb, -2, -1)
+
+
+class PositionalEncoding(nn.Module):
+    def __init__(self,
+                 dim: int,
+                 scale: float = math.pi * 2,
+                 temperature: float = 10000,
+                 normalize: bool = True,
+                 channel_last: bool = True,
+                 transpose_output: bool = False):
+        super().__init__()
+        dim = int(np.ceil(dim / 4) * 2)
+        self.dim = dim
+        inv_freq = 1.0 / (temperature**(torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
+        self.normalize = normalize
+        self.scale = scale
+        self.eps = 1e-6
+        self.channel_last = channel_last
+        self.transpose_output = transpose_output
+
+        self.cached_penc = None  # the cache is irrespective of the number of objects
+
+    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
+        """
+        :param tensor: A 4/5d tensor of size 
+            channel_last=True: (batch_size, h, w, c) or (batch_size, k, h, w, c)
+            channel_last=False: (batch_size, c, h, w) or (batch_size, k, c, h, w)
+        :return: positional encoding tensor that has the same shape as the input if the input is 4d
+                 if the input is 5d, the output is broadcastable along the k-dimension
+        """
+        if len(tensor.shape) != 4 and len(tensor.shape) != 5:
+            raise RuntimeError(f'The input tensor has to be 4/5d, got {tensor.shape}!')
+
+        if len(tensor.shape) == 5:
+            # take a sample from the k dimension
+            num_objects = tensor.shape[1]
+            tensor = tensor[:, 0]
+        else:
+            num_objects = None
+
+        if self.channel_last:
+            batch_size, h, w, c = tensor.shape
+        else:
+            batch_size, c, h, w = tensor.shape
+
+        if self.cached_penc is not None and self.cached_penc.shape == tensor.shape:
+            if num_objects is None:
+                return self.cached_penc
+            else:
+                return self.cached_penc.unsqueeze(1)
+
+        self.cached_penc = None
+
+        pos_y = torch.arange(h, device=tensor.device, dtype=self.inv_freq.dtype)
+        pos_x = torch.arange(w, device=tensor.device, dtype=self.inv_freq.dtype)
+        if self.normalize:
+            pos_y = pos_y / (pos_y[-1] + self.eps) * self.scale
+            pos_x = pos_x / (pos_x[-1] + self.eps) * self.scale
+
+        sin_inp_y = torch.einsum("i,j->ij", pos_y, self.inv_freq)
+        sin_inp_x = torch.einsum("i,j->ij", pos_x, self.inv_freq)
+        emb_y = get_emb(sin_inp_y).unsqueeze(1)
+        emb_x = get_emb(sin_inp_x)
+
+        emb = torch.zeros((h, w, self.dim * 2), device=tensor.device, dtype=tensor.dtype)
+        emb[:, :, :self.dim] = emb_x
+        emb[:, :, self.dim:] = emb_y
+
+        if not self.channel_last and self.transpose_output:
+            # cancelled out
+            pass
+        elif (not self.channel_last) or (self.transpose_output):
+            emb = emb.permute(2, 0, 1)
+
+        self.cached_penc = emb.unsqueeze(0).repeat(batch_size, 1, 1, 1)
+        if num_objects is None:
+            return self.cached_penc
+        else:
+            return self.cached_penc.unsqueeze(1)
+
+
+if __name__ == '__main__':
+    pe = PositionalEncoding(8).cuda()
+    input = torch.ones((1, 8, 8, 8)).cuda()
+    output = pe(input)
+    # print(output)
+    print(output[0, :, 0, 0])
+    print(output[0, :, 0, 5])
+    print(output[0, 0, :, 0])
+    print(output[0, 0, 0, :])

utility/matanyone2/matanyone/model/transformer/transformer_layers.py

@@ -0,0 +1,161 @@
+# Modified from PyTorch nn.Transformer
+
+from typing import List, Callable
+
+import torch
+from torch import Tensor
+import torch.nn as nn
+import torch.nn.functional as F
+from matanyone.model.channel_attn import CAResBlock
+
+
+class SelfAttention(nn.Module):
+    def __init__(self,
+                 dim: int,
+                 nhead: int,
+                 dropout: float = 0.0,
+                 batch_first: bool = True,
+                 add_pe_to_qkv: List[bool] = [True, True, False]):
+        super().__init__()
+        self.self_attn = nn.MultiheadAttention(dim, nhead, dropout=dropout, batch_first=batch_first)
+        self.norm = nn.LayerNorm(dim)
+        self.dropout = nn.Dropout(dropout)
+        self.add_pe_to_qkv = add_pe_to_qkv
+
+    def forward(self,
+                x: torch.Tensor,
+                pe: torch.Tensor,
+                attn_mask: bool = None,
+                key_padding_mask: bool = None) -> torch.Tensor:
+        x = self.norm(x)
+        if any(self.add_pe_to_qkv):
+            x_with_pe = x + pe
+            q = x_with_pe if self.add_pe_to_qkv[0] else x
+            k = x_with_pe if self.add_pe_to_qkv[1] else x
+            v = x_with_pe if self.add_pe_to_qkv[2] else x
+        else:
+            q = k = v = x
+
+        r = x
+        x = self.self_attn(q, k, v, attn_mask=attn_mask, key_padding_mask=key_padding_mask)[0]
+        return r + self.dropout(x)
+
+
+# https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html#torch.nn.functional.scaled_dot_product_attention
+class CrossAttention(nn.Module):
+    def __init__(self,
+                 dim: int,
+                 nhead: int,
+                 dropout: float = 0.0,
+                 batch_first: bool = True,
+                 add_pe_to_qkv: List[bool] = [True, True, False],
+                 residual: bool = True,
+                 norm: bool = True):
+        super().__init__()
+        self.cross_attn = nn.MultiheadAttention(dim,
+                                                nhead,
+                                                dropout=dropout,
+                                                batch_first=batch_first)
+        if norm:
+            self.norm = nn.LayerNorm(dim)
+        else:
+            self.norm = nn.Identity()
+        self.dropout = nn.Dropout(dropout)
+        self.add_pe_to_qkv = add_pe_to_qkv
+        self.residual = residual
+
+    def forward(self,
+                x: torch.Tensor,
+                mem: torch.Tensor,
+                x_pe: torch.Tensor,
+                mem_pe: torch.Tensor,
+                attn_mask: bool = None,
+                *,
+                need_weights: bool = False) -> (torch.Tensor, torch.Tensor):
+        x = self.norm(x)
+        if self.add_pe_to_qkv[0]:
+            q = x + x_pe
+        else:
+            q = x
+
+        if any(self.add_pe_to_qkv[1:]):
+            mem_with_pe = mem + mem_pe
+            k = mem_with_pe if self.add_pe_to_qkv[1] else mem
+            v = mem_with_pe if self.add_pe_to_qkv[2] else mem
+        else:
+            k = v = mem
+        r = x
+        x, weights = self.cross_attn(q,
+                                     k,
+                                     v,
+                                     attn_mask=attn_mask,
+                                     need_weights=need_weights,
+                                     average_attn_weights=False)
+
+        if self.residual:
+            return r + self.dropout(x), weights
+        else:
+            return self.dropout(x), weights
+
+
+class FFN(nn.Module):
+    def __init__(self, dim_in: int, dim_ff: int, activation=F.relu):
+        super().__init__()
+        self.linear1 = nn.Linear(dim_in, dim_ff)
+        self.linear2 = nn.Linear(dim_ff, dim_in)
+        self.norm = nn.LayerNorm(dim_in)
+
+        if isinstance(activation, str):
+            self.activation = _get_activation_fn(activation)
+        else:
+            self.activation = activation
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r = x
+        x = self.norm(x)
+        x = self.linear2(self.activation(self.linear1(x)))
+        x = r + x
+        return x
+
+
+class PixelFFN(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.dim = dim
+        self.conv = CAResBlock(dim, dim)
+
+    def forward(self, pixel: torch.Tensor, pixel_flat: torch.Tensor) -> torch.Tensor:
+        # pixel: batch_size * num_objects * dim * H * W
+        # pixel_flat: (batch_size*num_objects) * (H*W) * dim
+        bs, num_objects, _, h, w = pixel.shape
+        pixel_flat = pixel_flat.view(bs * num_objects, h, w, self.dim)
+        pixel_flat = pixel_flat.permute(0, 3, 1, 2).contiguous()
+
+        x = self.conv(pixel_flat)
+        x = x.view(bs, num_objects, self.dim, h, w)
+        return x
+
+
+class OutputFFN(nn.Module):
+    def __init__(self, dim_in: int, dim_out: int, activation=F.relu):
+        super().__init__()
+        self.linear1 = nn.Linear(dim_in, dim_out)
+        self.linear2 = nn.Linear(dim_out, dim_out)
+
+        if isinstance(activation, str):
+            self.activation = _get_activation_fn(activation)
+        else:
+            self.activation = activation
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.linear2(self.activation(self.linear1(x)))
+        return x
+
+
+def _get_activation_fn(activation: str) -> Callable[[Tensor], Tensor]:
+    if activation == "relu":
+        return F.relu
+    elif activation == "gelu":
+        return F.gelu
+
+    raise RuntimeError("activation should be relu/gelu, not {}".format(activation))

utility/matanyone2/matanyone/model/utils/memory_utils.py

@@ -0,0 +1,121 @@
+import math
+import torch
+from typing import Optional, Union, Tuple
+
+
+# @torch.jit.script
+def get_similarity(mk: torch.Tensor,
+                   ms: torch.Tensor,
+                   qk: torch.Tensor,
+                   qe: torch.Tensor,
+                   add_batch_dim: bool = False,
+                   uncert_mask = None) -> torch.Tensor:
+    # used for training/inference and memory reading/memory potentiation
+    # mk: B x CK x [N]    - Memory keys
+    # ms: B x  1 x [N]    - Memory shrinkage
+    # qk: B x CK x [HW/P] - Query keys
+    # qe: B x CK x [HW/P] - Query selection
+    # Dimensions in [] are flattened
+    # Return: B*N*HW
+    if add_batch_dim:
+        mk, ms = mk.unsqueeze(0), ms.unsqueeze(0)
+        qk, qe = qk.unsqueeze(0), qe.unsqueeze(0)
+
+    CK = mk.shape[1]
+
+    mk = mk.flatten(start_dim=2)
+    ms = ms.flatten(start_dim=1).unsqueeze(2) if ms is not None else None
+    qk = qk.flatten(start_dim=2)
+    qe = qe.flatten(start_dim=2) if qe is not None else None
+    
+    # query token selection based on temporal sparsity
+    if uncert_mask is not None:
+        uncert_mask = uncert_mask.flatten(start_dim=2)
+        uncert_mask = uncert_mask.expand(-1, 64, -1)
+        qk = qk * uncert_mask
+        qe = qe * uncert_mask
+    # Behold the work of DeeBeepMeep the Code Butcher !
+    if qe is not None:
+        # See XMem's appendix for derivation
+        mk = mk.transpose(1, 2)
+        a_sq = (mk.pow(2) @ qe)
+        two_ab =  mk @ (qk * qe)
+        two_ab *= 2
+        two_ab.sub_(a_sq)
+        del a_sq
+        b_sq = (qe * qk.pow(2)).sum(1, keepdim=True)
+        two_ab.sub_(b_sq)
+        similarity = two_ab
+        del b_sq, two_ab 
+        # similarity = (-a_sq + two_ab - b_sq)
+    else:
+        # similar to STCN if we don't have the selection term
+        a_sq = mk.pow(2).sum(1).unsqueeze(2)
+        two_ab = mk.transpose(1, 2) @ qk
+        two_ab *= 2
+        two_ab.sub_(a_sq)
+        del a_sq
+        similarity = two_ab
+        del two_ab 
+        # similarity = (-a_sq + two_ab)
+
+    if ms is not None:
+        similarity *= ms
+        similarity /=  math.sqrt(CK)
+        # similarity = similarity * ms / math.sqrt(CK)  # B*N*HW
+    else:
+        similarity /=  math.sqrt(CK)
+        # similarity = similarity / math.sqrt(CK)  # B*N*HW
+
+    return similarity
+
+
+def do_softmax(
+        similarity: torch.Tensor,
+        top_k: Optional[int] = None,
+        inplace: bool = False,
+        return_usage: bool = False) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+    # normalize similarity with top-k softmax
+    # similarity: B x N x [HW/P]
+    # use inplace with care
+    if top_k is not None:
+        values, indices = torch.topk(similarity, k=top_k, dim=1)
+
+        x_exp = values.exp_()
+        x_exp /= torch.sum(x_exp, dim=1, keepdim=True)
+        if inplace:
+            similarity.zero_().scatter_(1, indices, x_exp)  # B*N*HW
+            affinity = similarity
+        else:
+            affinity = torch.zeros_like(similarity).scatter_(1, indices, x_exp)  # B*N*HW
+    else:
+        maxes = torch.max(similarity, dim=1, keepdim=True)[0]
+        x_exp = torch.exp(similarity - maxes)
+        x_exp_sum = torch.sum(x_exp, dim=1, keepdim=True)
+        affinity = x_exp / x_exp_sum
+        indices = None
+
+    if return_usage:
+        return affinity, affinity.sum(dim=2)
+
+    return affinity
+
+
+def get_affinity(mk: torch.Tensor, ms: torch.Tensor, qk: torch.Tensor,
+                 qe: torch.Tensor, uncert_mask = None) -> torch.Tensor:
+    # shorthand used in training with no top-k
+    similarity = get_similarity(mk, ms, qk, qe, uncert_mask=uncert_mask)
+    affinity = do_softmax(similarity)
+    return affinity
+
+def readout(affinity: torch.Tensor, mv: torch.Tensor, uncert_mask: torch.Tensor=None) -> torch.Tensor:
+    B, CV, T, H, W = mv.shape
+
+    mo = mv.view(B, CV, T * H * W)
+    mem = torch.bmm(mo, affinity)
+    if uncert_mask is not None:
+        uncert_mask = uncert_mask.flatten(start_dim=2).expand(-1, CV, -1)
+        mem = mem * uncert_mask
+    mem = mem.view(B, CV, H, W)
+
+    return mem

utility/matanyone2/matanyone/model/utils/parameter_groups.py

@@ -0,0 +1,72 @@
+import logging
+
+log = logging.getLogger()
+
+
+def get_parameter_groups(model, stage_cfg, print_log=False):
+    """
+    Assign different weight decays and learning rates to different parameters.
+    Returns a parameter group which can be passed to the optimizer.
+    """
+    weight_decay = stage_cfg.weight_decay
+    embed_weight_decay = stage_cfg.embed_weight_decay
+    backbone_lr_ratio = stage_cfg.backbone_lr_ratio
+    base_lr = stage_cfg.learning_rate
+
+    backbone_params = []
+    embed_params = []
+    other_params = []
+
+    embedding_names = ['summary_pos', 'query_init', 'query_emb', 'obj_pe']
+    embedding_names = [e + '.weight' for e in embedding_names]
+
+    # inspired by detectron2
+    memo = set()
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+        # Avoid duplicating parameters
+        if param in memo:
+            continue
+        memo.add(param)
+
+        if name.startswith('module'):
+            name = name[7:]
+
+        inserted = False
+        if name.startswith('pixel_encoder.'):
+            backbone_params.append(param)
+            inserted = True
+            if print_log:
+                log.info(f'{name} counted as a backbone parameter.')
+        else:
+            for e in embedding_names:
+                if name.endswith(e):
+                    embed_params.append(param)
+                    inserted = True
+                    if print_log:
+                        log.info(f'{name} counted as an embedding parameter.')
+                    break
+
+        if not inserted:
+            other_params.append(param)
+
+    parameter_groups = [
+        {
+            'params': backbone_params,
+            'lr': base_lr * backbone_lr_ratio,
+            'weight_decay': weight_decay
+        },
+        {
+            'params': embed_params,
+            'lr': base_lr,
+            'weight_decay': embed_weight_decay
+        },
+        {
+            'params': other_params,
+            'lr': base_lr,
+            'weight_decay': weight_decay
+        },
+    ]
+
+    return parameter_groups

utility/matanyone2/matanyone/model/utils/resnet.py

@@ -0,0 +1,179 @@
+"""
+resnet.py - A modified ResNet structure
+We append extra channels to the first conv by some network surgery
+"""
+
+from collections import OrderedDict
+import math
+
+import torch
+import torch.nn as nn
+from torch.utils import model_zoo
+
+
+def load_weights_add_extra_dim(target, source_state, extra_dim=1):
+    new_dict = OrderedDict()
+
+    for k1, v1 in target.state_dict().items():
+        if not 'num_batches_tracked' in k1:
+            if k1 in source_state:
+                tar_v = source_state[k1]
+
+                if v1.shape != tar_v.shape:
+                    # Init the new segmentation channel with zeros
+                    # print(v1.shape, tar_v.shape)
+                    c, _, w, h = v1.shape
+                    pads = torch.zeros((c, extra_dim, w, h), device=tar_v.device)
+                    nn.init.orthogonal_(pads)
+                    tar_v = torch.cat([tar_v, pads], 1)
+
+                new_dict[k1] = tar_v
+
+    target.load_state_dict(new_dict)
+
+
+model_urls = {
+    'resnet18': 'http://download.pytorch.org/models/resnet18-5c106cde.pth',
+    'resnet50': 'http://download.pytorch.org/models/resnet50-19c8e357.pth',
+}
+
+
+def conv3x3(in_planes, out_planes, stride=1, dilation=1):
+    return nn.Conv2d(in_planes,
+                     out_planes,
+                     kernel_size=3,
+                     stride=stride,
+                     padding=dilation,
+                     dilation=dilation,
+                     bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, dilation=1):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride=stride, dilation=dilation)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes, stride=1, dilation=dilation)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, dilation=1):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes,
+                               planes,
+                               kernel_size=3,
+                               stride=stride,
+                               dilation=dilation,
+                               padding=dilation,
+                               bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * 4)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+    def __init__(self, block, layers=(3, 4, 23, 3), extra_dim=0):
+        self.inplanes = 64
+        super(ResNet, self).__init__()
+        self.conv1 = nn.Conv2d(3 + extra_dim, 64, kernel_size=7, stride=2, padding=3, bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilation=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes,
+                          planes * block.expansion,
+                          kernel_size=1,
+                          stride=stride,
+                          bias=False),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = [block(self.inplanes, planes, stride, downsample)]
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes, dilation=dilation))
+
+        return nn.Sequential(*layers)
+
+
+def resnet18(pretrained=True, extra_dim=0):
+    model = ResNet(BasicBlock, [2, 2, 2, 2], extra_dim)
+    if pretrained:
+        load_weights_add_extra_dim(model, model_zoo.load_url(model_urls['resnet18']), extra_dim)
+    return model
+
+
+def resnet50(pretrained=True, extra_dim=0):
+    model = ResNet(Bottleneck, [3, 4, 6, 3], extra_dim)
+    if pretrained:
+        load_weights_add_extra_dim(model, model_zoo.load_url(model_urls['resnet50']), extra_dim)
+    return model

utility/matanyone2/matanyone/utils/get_default_model.py

@@ -0,0 +1,32 @@
+"""
+A helper function to get a default model for quick testing
+"""
+from omegaconf import OmegaConf, open_dict
+#from hydra import compose, initialize
+
+import torch
+from matanyone.model.matanyone import MatAnyone
+
+def get_matanyone_model(ckpt_path, device=None) -> MatAnyone:
+    #initialize(version_base='1.3.2', config_path="../config", job_name="eval_our_config")
+    #cfg = compose(config_name="eval_matanyone_config")
+    cfg = OmegaConf.load("configs/eval_matanyone_config.yaml")
+    
+    #with open_dict(cfg):  # Allow modifying read-only configs
+    #cfg.weights = ckpt_path  # Update weights path
+    
+    with open_dict(cfg):
+        cfg['weights'] = ckpt_path
+    #print(OmegaConf.to_yaml(cfg))
+
+    # Load the network weights
+    if device is not None:
+        matanyone = MatAnyone(cfg, single_object=True).to(device).eval()
+        model_weights = torch.load(cfg.weights, map_location=device, weights_only=True)
+    else:  # if device is not specified, `.cuda()` by default
+        matanyone = MatAnyone(cfg, single_object=True).cuda().eval()
+        model_weights = torch.load(cfg.weights, weights_only=True)
+        
+    matanyone.load_weights(model_weights)
+
+    return matanyone

utility/matanyone2/matanyone/utils/inference_utils.py

@@ -0,0 +1,54 @@
+import os
+import cv2
+import random
+import numpy as np
+
+import torch
+import torchvision
+
+IMAGE_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.JPG', '.JPEG', '.PNG')
+VIDEO_EXTENSIONS = ('.mp4', '.mov', '.avi', '.MP4', '.MOV', '.AVI')
+
+def read_frame_from_videos(frame_root):
+    if frame_root.endswith(VIDEO_EXTENSIONS):  # Video file path
+        video_name = os.path.basename(frame_root)[:-4]
+        frames, _, info = torchvision.io.read_video(filename=frame_root, pts_unit='sec', output_format='TCHW') # RGB
+        fps = info['video_fps']
+    else:
+        video_name = os.path.basename(frame_root)
+        frames = []
+        fr_lst = sorted(os.listdir(frame_root))
+        for fr in fr_lst:
+            frame = cv2.imread(os.path.join(frame_root, fr))[...,[2,1,0]] # RGB, HWC
+            frames.append(frame)
+        fps = 24  # default
+        frames = torch.Tensor(np.array(frames)).permute(0, 3, 1, 2).contiguous() # TCHW
+    
+    length = frames.shape[0]
+
+    return frames, fps, length, video_name
+
+def get_video_paths(input_root):
+    video_paths = []
+    for root, _, files in os.walk(input_root):
+        for file in files:
+            if file.lower().endswith(VIDEO_EXTENSIONS):
+                video_paths.append(os.path.join(root, file))
+    return sorted(video_paths)
+
+def str_to_list(value):
+    return list(map(int, value.split(',')))
+
+def gen_dilate(alpha, min_kernel_size, max_kernel_size): 
+    kernel_size = random.randint(min_kernel_size, max_kernel_size)
+    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size,kernel_size))
+    fg_and_unknown = np.array(np.not_equal(alpha, 0).astype(np.float32))
+    dilate = cv2.dilate(fg_and_unknown, kernel, iterations=1)*255
+    return dilate.astype(np.float32)
+
+def gen_erosion(alpha, min_kernel_size, max_kernel_size): 
+    kernel_size = random.randint(min_kernel_size, max_kernel_size)
+    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size,kernel_size))
+    fg = np.array(np.equal(alpha, 255).astype(np.float32))
+    erode = cv2.erode(fg, kernel, iterations=1)*255
+    return erode.astype(np.float32)

utility/matanyone2/matanyone/utils/tensor_utils.py

@@ -0,0 +1,62 @@
+from typing import List, Iterable
+import torch
+import torch.nn.functional as F
+
+
+# STM
+def pad_divide_by(in_img: torch.Tensor, d: int) -> (torch.Tensor, Iterable[int]):
+    h, w = in_img.shape[-2:]
+
+    if h % d > 0:
+        new_h = h + d - h % d
+    else:
+        new_h = h
+    if w % d > 0:
+        new_w = w + d - w % d
+    else:
+        new_w = w
+    lh, uh = int((new_h - h) / 2), int(new_h - h) - int((new_h - h) / 2)
+    lw, uw = int((new_w - w) / 2), int(new_w - w) - int((new_w - w) / 2)
+    pad_array = (int(lw), int(uw), int(lh), int(uh))
+    out = F.pad(in_img, pad_array)
+    return out, pad_array
+
+
+def unpad(img: torch.Tensor, pad: Iterable[int]) -> torch.Tensor:
+    if len(img.shape) == 4:
+        if pad[2] + pad[3] > 0:
+            img = img[:, :, pad[2]:-pad[3], :]
+        if pad[0] + pad[1] > 0:
+            img = img[:, :, :, pad[0]:-pad[1]]
+    elif len(img.shape) == 3:
+        if pad[2] + pad[3] > 0:
+            img = img[:, pad[2]:-pad[3], :]
+        if pad[0] + pad[1] > 0:
+            img = img[:, :, pad[0]:-pad[1]]
+    elif len(img.shape) == 5:
+        if pad[2] + pad[3] > 0:
+            img = img[:, :, :, pad[2]:-pad[3], :]
+        if pad[0] + pad[1] > 0:
+            img = img[:, :, :, :, pad[0]:-pad[1]]
+    else:
+        raise NotImplementedError
+    return img
+
+
+# @torch.jit.script
+def aggregate(prob: torch.Tensor, dim: int) -> torch.Tensor:
+    with torch.amp.autocast("cuda",enabled=False):
+        prob = prob.float()
+        new_prob = torch.cat([torch.prod(1 - prob, dim=dim, keepdim=True), prob],
+                             dim).clamp(1e-7, 1 - 1e-7)
+        logits = torch.log((new_prob / (1 - new_prob)))  # (0, 1) --> (-inf, inf)
+
+        return logits
+
+
+# @torch.jit.script
+def cls_to_one_hot(cls_gt: torch.Tensor, num_objects: int) -> torch.Tensor:
+    # cls_gt: B*1*H*W
+    B, _, H, W = cls_gt.shape
+    one_hot = torch.zeros(B, num_objects + 1, H, W, device=cls_gt.device).scatter_(1, cls_gt, 1)
+    return one_hot

utility/sam/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from hydra import initialize_config_module
+from hydra.core.global_hydra import GlobalHydra
+
+if not GlobalHydra.instance().is_initialized():
+    initialize_config_module("sam2", version_base="1.2")

utility/sam/automatic_mask_generator.py

@@ -0,0 +1,454 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Adapted from https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/automatic_mask_generator.py
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+from torchvision.ops.boxes import batched_nms, box_area  # type: ignore
+
+from sam2.modeling.sam2_base import SAM2Base
+from sam2.sam2_image_predictor import SAM2ImagePredictor
+from sam2.utils.amg import (
+    area_from_rle,
+    batch_iterator,
+    batched_mask_to_box,
+    box_xyxy_to_xywh,
+    build_all_layer_point_grids,
+    calculate_stability_score,
+    coco_encode_rle,
+    generate_crop_boxes,
+    is_box_near_crop_edge,
+    mask_to_rle_pytorch,
+    MaskData,
+    remove_small_regions,
+    rle_to_mask,
+    uncrop_boxes_xyxy,
+    uncrop_masks,
+    uncrop_points,
+)
+
+
+class SAM2AutomaticMaskGenerator:
+    def __init__(
+        self,
+        model: SAM2Base,
+        points_per_side: Optional[int] = 32,
+        points_per_batch: int = 64,
+        pred_iou_thresh: float = 0.8,
+        stability_score_thresh: float = 0.95,
+        stability_score_offset: float = 1.0,
+        mask_threshold: float = 0.0,
+        box_nms_thresh: float = 0.7,
+        crop_n_layers: int = 0,
+        crop_nms_thresh: float = 0.7,
+        crop_overlap_ratio: float = 512 / 1500,
+        crop_n_points_downscale_factor: int = 1,
+        point_grids: Optional[List[np.ndarray]] = None,
+        min_mask_region_area: int = 0,
+        output_mode: str = "binary_mask",
+        use_m2m: bool = False,
+        multimask_output: bool = True,
+        **kwargs,
+    ) -> None:
+        """
+        Using a SAM 2 model, generates masks for the entire image.
+        Generates a grid of point prompts over the image, then filters
+        low quality and duplicate masks. The default settings are chosen
+        for SAM 2 with a HieraL backbone.
+
+        Arguments:
+          model (Sam): The SAM 2 model to use for mask prediction.
+          points_per_side (int or None): The number of points to be sampled
+            along one side of the image. The total number of points is
+            points_per_side**2. If None, 'point_grids' must provide explicit
+            point sampling.
+          points_per_batch (int): Sets the number of points run simultaneously
+            by the model. Higher numbers may be faster but use more GPU memory.
+          pred_iou_thresh (float): A filtering threshold in [0,1], using the
+            model's predicted mask quality.
+          stability_score_thresh (float): A filtering threshold in [0,1], using
+            the stability of the mask under changes to the cutoff used to binarize
+            the model's mask predictions.
+          stability_score_offset (float): The amount to shift the cutoff when
+            calculated the stability score.
+          mask_threshold (float): Threshold for binarizing the mask logits
+          box_nms_thresh (float): The box IoU cutoff used by non-maximal
+            suppression to filter duplicate masks.
+          crop_n_layers (int): If >0, mask prediction will be run again on
+            crops of the image. Sets the number of layers to run, where each
+            layer has 2**i_layer number of image crops.
+          crop_nms_thresh (float): The box IoU cutoff used by non-maximal
+            suppression to filter duplicate masks between different crops.
+          crop_overlap_ratio (float): Sets the degree to which crops overlap.
+            In the first crop layer, crops will overlap by this fraction of
+            the image length. Later layers with more crops scale down this overlap.
+          crop_n_points_downscale_factor (int): The number of points-per-side
+            sampled in layer n is scaled down by crop_n_points_downscale_factor**n.
+          point_grids (list(np.ndarray) or None): A list over explicit grids
+            of points used for sampling, normalized to [0,1]. The nth grid in the
+            list is used in the nth crop layer. Exclusive with points_per_side.
+          min_mask_region_area (int): If >0, postprocessing will be applied
+            to remove disconnected regions and holes in masks with area smaller
+            than min_mask_region_area. Requires opencv.
+          output_mode (str): The form masks are returned in. Can be 'binary_mask',
+            'uncompressed_rle', or 'coco_rle'. 'coco_rle' requires pycocotools.
+            For large resolutions, 'binary_mask' may consume large amounts of
+            memory.
+          use_m2m (bool): Whether to add a one step refinement using previous mask predictions.
+          multimask_output (bool): Whether to output multimask at each point of the grid.
+        """
+
+        assert (points_per_side is None) != (
+            point_grids is None
+        ), "Exactly one of points_per_side or point_grid must be provided."
+        if points_per_side is not None:
+            self.point_grids = build_all_layer_point_grids(
+                points_per_side,
+                crop_n_layers,
+                crop_n_points_downscale_factor,
+            )
+        elif point_grids is not None:
+            self.point_grids = point_grids
+        else:
+            raise ValueError("Can't have both points_per_side and point_grid be None.")
+
+        assert output_mode in [
+            "binary_mask",
+            "uncompressed_rle",
+            "coco_rle",
+        ], f"Unknown output_mode {output_mode}."
+        if output_mode == "coco_rle":
+            try:
+                from pycocotools import mask as mask_utils  # type: ignore  # noqa: F401
+            except ImportError as e:
+                print("Please install pycocotools")
+                raise e
+
+        self.predictor = SAM2ImagePredictor(
+            model,
+            max_hole_area=min_mask_region_area,
+            max_sprinkle_area=min_mask_region_area,
+        )
+        self.points_per_batch = points_per_batch
+        self.pred_iou_thresh = pred_iou_thresh
+        self.stability_score_thresh = stability_score_thresh
+        self.stability_score_offset = stability_score_offset
+        self.mask_threshold = mask_threshold
+        self.box_nms_thresh = box_nms_thresh
+        self.crop_n_layers = crop_n_layers
+        self.crop_nms_thresh = crop_nms_thresh
+        self.crop_overlap_ratio = crop_overlap_ratio
+        self.crop_n_points_downscale_factor = crop_n_points_downscale_factor
+        self.min_mask_region_area = min_mask_region_area
+        self.output_mode = output_mode
+        self.use_m2m = use_m2m
+        self.multimask_output = multimask_output
+
+    @classmethod
+    def from_pretrained(cls, model_id: str, **kwargs) -> "SAM2AutomaticMaskGenerator":
+        """
+        Load a pretrained model from the Hugging Face hub.
+
+        Arguments:
+          model_id (str): The Hugging Face repository ID.
+          **kwargs: Additional arguments to pass to the model constructor.
+
+        Returns:
+          (SAM2AutomaticMaskGenerator): The loaded model.
+        """
+        from sam2.build_sam import build_sam2_hf
+
+        sam_model = build_sam2_hf(model_id, **kwargs)
+        return cls(sam_model, **kwargs)
+
+    @torch.no_grad()
+    def generate(self, image: np.ndarray) -> List[Dict[str, Any]]:
+        """
+        Generates masks for the given image.
+
+        Arguments:
+          image (np.ndarray): The image to generate masks for, in HWC uint8 format.
+
+        Returns:
+           list(dict(str, any)): A list over records for masks. Each record is
+             a dict containing the following keys:
+               segmentation (dict(str, any) or np.ndarray): The mask. If
+                 output_mode='binary_mask', is an array of shape HW. Otherwise,
+                 is a dictionary containing the RLE.
+               bbox (list(float)): The box around the mask, in XYWH format.
+               area (int): The area in pixels of the mask.
+               predicted_iou (float): The model's own prediction of the mask's
+                 quality. This is filtered by the pred_iou_thresh parameter.
+               point_coords (list(list(float))): The point coordinates input
+                 to the model to generate this mask.
+               stability_score (float): A measure of the mask's quality. This
+                 is filtered on using the stability_score_thresh parameter.
+               crop_box (list(float)): The crop of the image used to generate
+                 the mask, given in XYWH format.
+        """
+
+        # Generate masks
+        mask_data = self._generate_masks(image)
+
+        # Encode masks
+        if self.output_mode == "coco_rle":
+            mask_data["segmentations"] = [
+                coco_encode_rle(rle) for rle in mask_data["rles"]
+            ]
+        elif self.output_mode == "binary_mask":
+            mask_data["segmentations"] = [rle_to_mask(rle) for rle in mask_data["rles"]]
+        else:
+            mask_data["segmentations"] = mask_data["rles"]
+
+        # Write mask records
+        curr_anns = []
+        for idx in range(len(mask_data["segmentations"])):
+            ann = {
+                "segmentation": mask_data["segmentations"][idx],
+                "area": area_from_rle(mask_data["rles"][idx]),
+                "bbox": box_xyxy_to_xywh(mask_data["boxes"][idx]).tolist(),
+                "predicted_iou": mask_data["iou_preds"][idx].item(),
+                "point_coords": [mask_data["points"][idx].tolist()],
+                "stability_score": mask_data["stability_score"][idx].item(),
+                "crop_box": box_xyxy_to_xywh(mask_data["crop_boxes"][idx]).tolist(),
+            }
+            curr_anns.append(ann)
+
+        return curr_anns
+
+    def _generate_masks(self, image: np.ndarray) -> MaskData:
+        orig_size = image.shape[:2]
+        crop_boxes, layer_idxs = generate_crop_boxes(
+            orig_size, self.crop_n_layers, self.crop_overlap_ratio
+        )
+
+        # Iterate over image crops
+        data = MaskData()
+        for crop_box, layer_idx in zip(crop_boxes, layer_idxs):
+            crop_data = self._process_crop(image, crop_box, layer_idx, orig_size)
+            data.cat(crop_data)
+
+        # Remove duplicate masks between crops
+        if len(crop_boxes) > 1:
+            # Prefer masks from smaller crops
+            scores = 1 / box_area(data["crop_boxes"])
+            scores = scores.to(data["boxes"].device)
+            keep_by_nms = batched_nms(
+                data["boxes"].float(),
+                scores,
+                torch.zeros_like(data["boxes"][:, 0]),  # categories
+                iou_threshold=self.crop_nms_thresh,
+            )
+            data.filter(keep_by_nms)
+        data.to_numpy()
+        return data
+
+    def _process_crop(
+        self,
+        image: np.ndarray,
+        crop_box: List[int],
+        crop_layer_idx: int,
+        orig_size: Tuple[int, ...],
+    ) -> MaskData:
+        # Crop the image and calculate embeddings
+        x0, y0, x1, y1 = crop_box
+        cropped_im = image[y0:y1, x0:x1, :]
+        cropped_im_size = cropped_im.shape[:2]
+        self.predictor.set_image(cropped_im)
+
+        # Get points for this crop
+        points_scale = np.array(cropped_im_size)[None, ::-1]
+        points_for_image = self.point_grids[crop_layer_idx] * points_scale
+
+        # Generate masks for this crop in batches
+        data = MaskData()
+        for (points,) in batch_iterator(self.points_per_batch, points_for_image):
+            batch_data = self._process_batch(
+                points, cropped_im_size, crop_box, orig_size, normalize=True
+            )
+            data.cat(batch_data)
+            del batch_data
+        self.predictor.reset_predictor()
+
+        # Remove duplicates within this crop.
+        keep_by_nms = batched_nms(
+            data["boxes"].float(),
+            data["iou_preds"],
+            torch.zeros_like(data["boxes"][:, 0]),  # categories
+            iou_threshold=self.box_nms_thresh,
+        )
+        data.filter(keep_by_nms)
+
+        # Return to the original image frame
+        data["boxes"] = uncrop_boxes_xyxy(data["boxes"], crop_box)
+        data["points"] = uncrop_points(data["points"], crop_box)
+        data["crop_boxes"] = torch.tensor([crop_box for _ in range(len(data["rles"]))])
+
+        return data
+
+    def _process_batch(
+        self,
+        points: np.ndarray,
+        im_size: Tuple[int, ...],
+        crop_box: List[int],
+        orig_size: Tuple[int, ...],
+        normalize=False,
+    ) -> MaskData:
+        orig_h, orig_w = orig_size
+
+        # Run model on this batch
+        points = torch.as_tensor(
+            points, dtype=torch.float32, device=self.predictor.device
+        )
+        in_points = self.predictor._transforms.transform_coords(
+            points, normalize=normalize, orig_hw=im_size
+        )
+        in_labels = torch.ones(
+            in_points.shape[0], dtype=torch.int, device=in_points.device
+        )
+        masks, iou_preds, low_res_masks = self.predictor._predict(
+            in_points[:, None, :],
+            in_labels[:, None],
+            multimask_output=self.multimask_output,
+            return_logits=True,
+        )
+
+        # Serialize predictions and store in MaskData
+        data = MaskData(
+            masks=masks.flatten(0, 1),
+            iou_preds=iou_preds.flatten(0, 1),
+            points=points.repeat_interleave(masks.shape[1], dim=0),
+            low_res_masks=low_res_masks.flatten(0, 1),
+        )
+        del masks
+
+        if not self.use_m2m:
+            # Filter by predicted IoU
+            if self.pred_iou_thresh > 0.0:
+                keep_mask = data["iou_preds"] > self.pred_iou_thresh
+                data.filter(keep_mask)
+
+            # Calculate and filter by stability score
+            data["stability_score"] = calculate_stability_score(
+                data["masks"], self.mask_threshold, self.stability_score_offset
+            )
+            if self.stability_score_thresh > 0.0:
+                keep_mask = data["stability_score"] >= self.stability_score_thresh
+                data.filter(keep_mask)
+        else:
+            # One step refinement using previous mask predictions
+            in_points = self.predictor._transforms.transform_coords(
+                data["points"], normalize=normalize, orig_hw=im_size
+            )
+            labels = torch.ones(
+                in_points.shape[0], dtype=torch.int, device=in_points.device
+            )
+            masks, ious = self.refine_with_m2m(
+                in_points, labels, data["low_res_masks"], self.points_per_batch
+            )
+            data["masks"] = masks.squeeze(1)
+            data["iou_preds"] = ious.squeeze(1)
+
+            if self.pred_iou_thresh > 0.0:
+                keep_mask = data["iou_preds"] > self.pred_iou_thresh
+                data.filter(keep_mask)
+
+            data["stability_score"] = calculate_stability_score(
+                data["masks"], self.mask_threshold, self.stability_score_offset
+            )
+            if self.stability_score_thresh > 0.0:
+                keep_mask = data["stability_score"] >= self.stability_score_thresh
+                data.filter(keep_mask)
+
+        # Threshold masks and calculate boxes
+        data["masks"] = data["masks"] > self.mask_threshold
+        data["boxes"] = batched_mask_to_box(data["masks"])
+
+        # Filter boxes that touch crop boundaries
+        keep_mask = ~is_box_near_crop_edge(
+            data["boxes"], crop_box, [0, 0, orig_w, orig_h]
+        )
+        if not torch.all(keep_mask):
+            data.filter(keep_mask)
+
+        # Compress to RLE
+        data["masks"] = uncrop_masks(data["masks"], crop_box, orig_h, orig_w)
+        data["rles"] = mask_to_rle_pytorch(data["masks"])
+        del data["masks"]
+
+        return data
+
+    @staticmethod
+    def postprocess_small_regions(
+        mask_data: MaskData, min_area: int, nms_thresh: float
+    ) -> MaskData:
+        """
+        Removes small disconnected regions and holes in masks, then reruns
+        box NMS to remove any new duplicates.
+
+        Edits mask_data in place.
+
+        Requires open-cv as a dependency.
+        """
+        if len(mask_data["rles"]) == 0:
+            return mask_data
+
+        # Filter small disconnected regions and holes
+        new_masks = []
+        scores = []
+        for rle in mask_data["rles"]:
+            mask = rle_to_mask(rle)
+
+            mask, changed = remove_small_regions(mask, min_area, mode="holes")
+            unchanged = not changed
+            mask, changed = remove_small_regions(mask, min_area, mode="islands")
+            unchanged = unchanged and not changed
+
+            new_masks.append(torch.as_tensor(mask).unsqueeze(0))
+            # Give score=0 to changed masks and score=1 to unchanged masks
+            # so NMS will prefer ones that didn't need postprocessing
+            scores.append(float(unchanged))
+
+        # Recalculate boxes and remove any new duplicates
+        masks = torch.cat(new_masks, dim=0)
+        boxes = batched_mask_to_box(masks)
+        keep_by_nms = batched_nms(
+            boxes.float(),
+            torch.as_tensor(scores),
+            torch.zeros_like(boxes[:, 0]),  # categories
+            iou_threshold=nms_thresh,
+        )
+
+        # Only recalculate RLEs for masks that have changed
+        for i_mask in keep_by_nms:
+            if scores[i_mask] == 0.0:
+                mask_torch = masks[i_mask].unsqueeze(0)
+                mask_data["rles"][i_mask] = mask_to_rle_pytorch(mask_torch)[0]
+                mask_data["boxes"][i_mask] = boxes[i_mask]  # update res directly
+        mask_data.filter(keep_by_nms)
+
+        return mask_data
+
+    def refine_with_m2m(self, points, point_labels, low_res_masks, points_per_batch):
+        new_masks = []
+        new_iou_preds = []
+
+        for cur_points, cur_point_labels, low_res_mask in batch_iterator(
+            points_per_batch, points, point_labels, low_res_masks
+        ):
+            best_masks, best_iou_preds, _ = self.predictor._predict(
+                cur_points[:, None, :],
+                cur_point_labels[:, None],
+                mask_input=low_res_mask[:, None, :],
+                multimask_output=False,
+                return_logits=True,
+            )
+            new_masks.append(best_masks)
+            new_iou_preds.append(best_iou_preds)
+        masks = torch.cat(new_masks, dim=0)
+        return masks, torch.cat(new_iou_preds, dim=0)

utility/sam/build_sam.py

@@ -0,0 +1,173 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+import os
+
+import torch
+from hydra import compose
+from hydra.utils import instantiate
+from omegaconf import OmegaConf
+
+import sam2
+
+# Check if the user is running Python from the parent directory of the sam2 repo
+# (i.e. the directory where this repo is cloned into) -- this is not supported since
+# it could shadow the sam2 package and cause issues.
+if os.path.isdir(os.path.join(sam2.__path__[0], "sam2")):
+    # If the user has "sam2/sam2" in their path, they are likey importing the repo itself
+    # as "sam2" rather than importing the "sam2" python package (i.e. "sam2/sam2" directory).
+    # This typically happens because the user is running Python from the parent directory
+    # that contains the sam2 repo they cloned.
+    raise RuntimeError(
+        "You're likely running Python from the parent directory of the sam2 repository "
+        "(i.e. the directory where https://github.com/facebookresearch/sam2 is cloned into). "
+        "This is not supported since the `sam2` Python package could be shadowed by the "
+        "repository name (the repository is also named `sam2` and contains the Python package "
+        "in `sam2/sam2`). Please run Python from another directory (e.g. from the repo dir "
+        "rather than its parent dir, or from your home directory) after installing SAM 2."
+    )
+
+
+HF_MODEL_ID_TO_FILENAMES = {
+    "facebook/sam2-hiera-tiny": (
+        "configs/sam2/sam2_hiera_t.yaml",
+        "sam2_hiera_tiny.pt",
+    ),
+    "facebook/sam2-hiera-small": (
+        "configs/sam2/sam2_hiera_s.yaml",
+        "sam2_hiera_small.pt",
+    ),
+    "facebook/sam2-hiera-base-plus": (
+        "configs/sam2/sam2_hiera_b+.yaml",
+        "sam2_hiera_base_plus.pt",
+    ),
+    "facebook/sam2-hiera-large": (
+        "configs/sam2/sam2_hiera_l.yaml",
+        "sam2_hiera_large.pt",
+    ),
+    "facebook/sam2.1-hiera-tiny": (
+        "configs/sam2.1/sam2.1_hiera_t.yaml",
+        "sam2.1_hiera_tiny.pt",
+    ),
+    "facebook/sam2.1-hiera-small": (
+        "configs/sam2.1/sam2.1_hiera_s.yaml",
+        "sam2.1_hiera_small.pt",
+    ),
+    "facebook/sam2.1-hiera-base-plus": (
+        "configs/sam2.1/sam2.1_hiera_b+.yaml",
+        "sam2.1_hiera_base_plus.pt",
+    ),
+    "facebook/sam2.1-hiera-large": (
+        "configs/sam2.1/sam2.1_hiera_l.yaml",
+        "sam2.1_hiera_large.pt",
+    ),
+}
+
+
+def build_sam2(
+    config_file,
+    ckpt_path=None,
+    device="cuda",
+    mode="eval",
+    hydra_overrides_extra=[],
+    apply_postprocessing=True,
+    **kwargs,
+):
+
+    if apply_postprocessing:
+        hydra_overrides_extra = hydra_overrides_extra.copy()
+        hydra_overrides_extra += [
+            # dynamically fall back to multi-mask if the single mask is not stable
+            "++model.sam_mask_decoder_extra_args.dynamic_multimask_via_stability=true",
+            "++model.sam_mask_decoder_extra_args.dynamic_multimask_stability_delta=0.05",
+            "++model.sam_mask_decoder_extra_args.dynamic_multimask_stability_thresh=0.98",
+        ]
+    # Read config and init model
+    cfg = compose(config_name=config_file, overrides=hydra_overrides_extra)
+    OmegaConf.resolve(cfg)
+    model = instantiate(cfg.model, _recursive_=True)
+    _load_checkpoint(model, ckpt_path)
+    model = model.to(device)
+    if mode == "eval":
+        model.eval()
+    return model
+
+
+def build_sam2_video_predictor(
+    config_file,
+    ckpt_path=None,
+    device="cuda",
+    mode="eval",
+    hydra_overrides_extra=[],
+    apply_postprocessing=True,
+    vos_optimized=False,
+    **kwargs,
+):
+    hydra_overrides = [
+        "++model._target_=sam2.sam2_video_predictor.SAM2VideoPredictor",
+    ]
+    if vos_optimized:
+        hydra_overrides = [
+            "++model._target_=sam2.sam2_video_predictor.SAM2VideoPredictorVOS",
+            "++model.compile_image_encoder=True",  # Let sam2_base handle this
+        ]
+
+    if apply_postprocessing:
+        hydra_overrides_extra = hydra_overrides_extra.copy()
+        hydra_overrides_extra += [
+            # dynamically fall back to multi-mask if the single mask is not stable
+            "++model.sam_mask_decoder_extra_args.dynamic_multimask_via_stability=true",
+            "++model.sam_mask_decoder_extra_args.dynamic_multimask_stability_delta=0.05",
+            "++model.sam_mask_decoder_extra_args.dynamic_multimask_stability_thresh=0.98",
+            # the sigmoid mask logits on interacted frames with clicks in the memory encoder so that the encoded masks are exactly as what users see from clicking
+            "++model.binarize_mask_from_pts_for_mem_enc=true",
+            # fill small holes in the low-res masks up to `fill_hole_area` (before resizing them to the original video resolution)
+            "++model.fill_hole_area=8",
+        ]
+    hydra_overrides.extend(hydra_overrides_extra)
+
+    # Read config and init model
+    cfg = compose(config_name=config_file, overrides=hydra_overrides)
+    OmegaConf.resolve(cfg)
+    model = instantiate(cfg.model, _recursive_=True)
+    _load_checkpoint(model, ckpt_path)
+    model = model.to(device)
+    if mode == "eval":
+        model.eval()
+    return model
+
+
+def _hf_download(model_id):
+    from huggingface_hub import hf_hub_download
+
+    config_name, checkpoint_name = HF_MODEL_ID_TO_FILENAMES[model_id]
+    ckpt_path = hf_hub_download(repo_id=model_id, filename=checkpoint_name)
+    return config_name, ckpt_path
+
+
+def build_sam2_hf(model_id, **kwargs):
+    config_name, ckpt_path = _hf_download(model_id)
+    return build_sam2(config_file=config_name, ckpt_path=ckpt_path, **kwargs)
+
+
+def build_sam2_video_predictor_hf(model_id, **kwargs):
+    config_name, ckpt_path = _hf_download(model_id)
+    return build_sam2_video_predictor(
+        config_file=config_name, ckpt_path=ckpt_path, **kwargs
+    )
+
+
+def _load_checkpoint(model, ckpt_path):
+    if ckpt_path is not None:
+        sd = torch.load(ckpt_path, map_location="cpu", weights_only=True)["model"]
+        missing_keys, unexpected_keys = model.load_state_dict(sd)
+        if missing_keys:
+            logging.error(missing_keys)
+            raise RuntimeError()
+        if unexpected_keys:
+            logging.error(unexpected_keys)
+            raise RuntimeError()

utility/sam/configs/sam2.1/efficienttam_s.yaml

@@ -0,0 +1,123 @@
+# @package _global_
+
+# Model
+model:
+  _target_: sam2.modeling.sam2_base.SAM2Base
+  image_encoder:
+    _target_: sam2.modeling.backbones.image_encoder.ImageEncoder
+    scalp: 0
+    trunk:
+      _target_: sam2.modeling.backbones.vitdet.ViT
+      patch_size: 16
+      embed_dim: 384
+      depth: 12
+      num_heads: 6
+      mlp_ratio: 4.0
+      qkv_bias: true
+      drop_path_rate: 0.0
+      use_rel_pos: false
+      window_size: 14
+      window_block_indexes: [0, 1, 3, 4, 6, 7, 9, 10]
+    neck:
+      _target_: sam2.modeling.backbones.image_encoder.ViTDetNeck
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 256
+        normalize: true
+        scale: null
+        temperature: 10000
+      d_model: 256
+      backbone_channel_list: [384,]
+      neck_norm: LN
+
+  memory_attention:
+    _target_: sam2.modeling.memory_attention.MemoryAttention
+    d_model: 256
+    pos_enc_at_input: true
+    layer:
+      _target_: sam2.modeling.memory_attention.MemoryAttentionLayer
+      activation: relu
+      dim_feedforward: 2048
+      dropout: 0.1
+      pos_enc_at_attn: false
+      self_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [32, 32]
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+      d_model: 256
+      pos_enc_at_cross_attn_keys: true
+      pos_enc_at_cross_attn_queries: false
+      cross_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [32, 32]
+        rope_k_repeat: True
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+        kv_in_dim: 64
+    num_layers: 4
+
+  memory_encoder:
+      _target_: sam2.modeling.memory_encoder.MemoryEncoder
+      out_dim: 64
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 64
+        normalize: true
+        scale: null
+        temperature: 10000
+      mask_downsampler:
+        _target_: sam2.modeling.memory_encoder.MaskDownSampler
+        kernel_size: 3
+        stride: 2
+        padding: 1
+      fuser:
+        _target_: sam2.modeling.memory_encoder.Fuser
+        layer:
+          _target_: sam2.modeling.memory_encoder.CXBlock
+          dim: 256
+          kernel_size: 7
+          padding: 3
+          layer_scale_init_value: 1e-6
+          use_dwconv: True  # depth-wise convs
+        num_layers: 2
+
+  num_maskmem: 7
+  image_size: 1024
+  # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask
+  # SAM decoder
+  sigmoid_scale_for_mem_enc: 20.0
+  sigmoid_bias_for_mem_enc: -10.0
+  use_mask_input_as_output_without_sam: true
+  # Memory
+  directly_add_no_mem_embed: true
+  # use high-resolution feature map in the SAM mask decoder
+  # use_high_res_features_in_sam: true
+  use_high_res_features_in_sam: false
+  # output 3 masks on the first click on initial conditioning frames
+  multimask_output_in_sam: true
+  # SAM heads
+  iou_prediction_use_sigmoid: True
+  # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder
+  use_obj_ptrs_in_encoder: true
+  add_tpos_enc_to_obj_ptrs: false
+  only_obj_ptrs_in_the_past_for_eval: true
+  # object occlusion prediction
+  pred_obj_scores: true
+  pred_obj_scores_mlp: true
+  fixed_no_obj_ptr: true
+  # multimask tracking settings
+  multimask_output_for_tracking: true
+  use_multimask_token_for_obj_ptr: true
+  multimask_min_pt_num: 0
+  multimask_max_pt_num: 1
+  use_mlp_for_obj_ptr_proj: true
+  # Compilation flag
+  # HieraT does not currently support compilation, should always be set to False
+  compile_image_encoder: false

utility/sam/configs/sam2.1/sam2.1_hiera_b+.yaml

@@ -0,0 +1,116 @@
+# @package _global_
+
+# Model
+model:
+  _target_: sam2.modeling.sam2_base.SAM2Base
+  image_encoder:
+    _target_: sam2.modeling.backbones.image_encoder.ImageEncoder
+    scalp: 1
+    trunk:
+      _target_: sam2.modeling.backbones.hieradet.Hiera
+      embed_dim: 112
+      num_heads: 2
+    neck:
+      _target_: sam2.modeling.backbones.image_encoder.FpnNeck
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 256
+        normalize: true
+        scale: null
+        temperature: 10000
+      d_model: 256
+      backbone_channel_list: [896, 448, 224, 112]
+      fpn_top_down_levels: [2, 3]  # output level 0 and 1 directly use the backbone features
+      fpn_interp_model: nearest
+
+  memory_attention:
+    _target_: sam2.modeling.memory_attention.MemoryAttention
+    d_model: 256
+    pos_enc_at_input: true
+    layer:
+      _target_: sam2.modeling.memory_attention.MemoryAttentionLayer
+      activation: relu
+      dim_feedforward: 2048
+      dropout: 0.1
+      pos_enc_at_attn: false
+      self_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [64, 64]
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+      d_model: 256
+      pos_enc_at_cross_attn_keys: true
+      pos_enc_at_cross_attn_queries: false
+      cross_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [64, 64]
+        rope_k_repeat: True
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+        kv_in_dim: 64
+    num_layers: 4
+
+  memory_encoder:
+      _target_: sam2.modeling.memory_encoder.MemoryEncoder
+      out_dim: 64
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 64
+        normalize: true
+        scale: null
+        temperature: 10000
+      mask_downsampler:
+        _target_: sam2.modeling.memory_encoder.MaskDownSampler
+        kernel_size: 3
+        stride: 2
+        padding: 1
+      fuser:
+        _target_: sam2.modeling.memory_encoder.Fuser
+        layer:
+          _target_: sam2.modeling.memory_encoder.CXBlock
+          dim: 256
+          kernel_size: 7
+          padding: 3
+          layer_scale_init_value: 1e-6
+          use_dwconv: True  # depth-wise convs
+        num_layers: 2
+
+  num_maskmem: 7
+  image_size: 1024
+  # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask
+  sigmoid_scale_for_mem_enc: 20.0
+  sigmoid_bias_for_mem_enc: -10.0
+  use_mask_input_as_output_without_sam: true
+  # Memory
+  directly_add_no_mem_embed: true
+  no_obj_embed_spatial: true
+  # use high-resolution feature map in the SAM mask decoder
+  use_high_res_features_in_sam: true
+  # output 3 masks on the first click on initial conditioning frames
+  multimask_output_in_sam: true
+  # SAM heads
+  iou_prediction_use_sigmoid: True
+  # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder
+  use_obj_ptrs_in_encoder: true
+  add_tpos_enc_to_obj_ptrs: true
+  proj_tpos_enc_in_obj_ptrs: true
+  use_signed_tpos_enc_to_obj_ptrs: true
+  only_obj_ptrs_in_the_past_for_eval: true
+  # object occlusion prediction
+  pred_obj_scores: true
+  pred_obj_scores_mlp: true
+  fixed_no_obj_ptr: true
+  # multimask tracking settings
+  multimask_output_for_tracking: true
+  use_multimask_token_for_obj_ptr: true
+  multimask_min_pt_num: 0
+  multimask_max_pt_num: 1
+  use_mlp_for_obj_ptr_proj: true
+  # Compilation flag
+  compile_image_encoder: False

utility/sam/configs/sam2.1/sam2.1_hiera_l.yaml

@@ -0,0 +1,120 @@
+# @package _global_
+
+# Model
+model:
+  _target_: sam2.modeling.sam2_base.SAM2Base
+  image_encoder:
+    _target_: sam2.modeling.backbones.image_encoder.ImageEncoder
+    scalp: 1
+    trunk:
+      _target_: sam2.modeling.backbones.hieradet.Hiera
+      embed_dim: 144
+      num_heads: 2
+      stages: [2, 6, 36, 4]
+      global_att_blocks: [23, 33, 43]
+      window_pos_embed_bkg_spatial_size: [7, 7]
+      window_spec: [8, 4, 16, 8]
+    neck:
+      _target_: sam2.modeling.backbones.image_encoder.FpnNeck
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 256
+        normalize: true
+        scale: null
+        temperature: 10000
+      d_model: 256
+      backbone_channel_list: [1152, 576, 288, 144]
+      fpn_top_down_levels: [2, 3]  # output level 0 and 1 directly use the backbone features
+      fpn_interp_model: nearest
+
+  memory_attention:
+    _target_: sam2.modeling.memory_attention.MemoryAttention
+    d_model: 256
+    pos_enc_at_input: true
+    layer:
+      _target_: sam2.modeling.memory_attention.MemoryAttentionLayer
+      activation: relu
+      dim_feedforward: 2048
+      dropout: 0.1
+      pos_enc_at_attn: false
+      self_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [64, 64]
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+      d_model: 256
+      pos_enc_at_cross_attn_keys: true
+      pos_enc_at_cross_attn_queries: false
+      cross_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [64, 64]
+        rope_k_repeat: True
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+        kv_in_dim: 64
+    num_layers: 4
+
+  memory_encoder:
+      _target_: sam2.modeling.memory_encoder.MemoryEncoder
+      out_dim: 64
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 64
+        normalize: true
+        scale: null
+        temperature: 10000
+      mask_downsampler:
+        _target_: sam2.modeling.memory_encoder.MaskDownSampler
+        kernel_size: 3
+        stride: 2
+        padding: 1
+      fuser:
+        _target_: sam2.modeling.memory_encoder.Fuser
+        layer:
+          _target_: sam2.modeling.memory_encoder.CXBlock
+          dim: 256
+          kernel_size: 7
+          padding: 3
+          layer_scale_init_value: 1e-6
+          use_dwconv: True  # depth-wise convs
+        num_layers: 2
+
+  num_maskmem: 7
+  image_size: 1024
+  # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask
+  sigmoid_scale_for_mem_enc: 20.0
+  sigmoid_bias_for_mem_enc: -10.0
+  use_mask_input_as_output_without_sam: true
+  # Memory
+  directly_add_no_mem_embed: true
+  no_obj_embed_spatial: true
+  # use high-resolution feature map in the SAM mask decoder
+  use_high_res_features_in_sam: true
+  # output 3 masks on the first click on initial conditioning frames
+  multimask_output_in_sam: true
+  # SAM heads
+  iou_prediction_use_sigmoid: True
+  # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder
+  use_obj_ptrs_in_encoder: true
+  add_tpos_enc_to_obj_ptrs: true
+  proj_tpos_enc_in_obj_ptrs: true
+  use_signed_tpos_enc_to_obj_ptrs: true
+  only_obj_ptrs_in_the_past_for_eval: true
+  # object occlusion prediction
+  pred_obj_scores: true
+  pred_obj_scores_mlp: true
+  fixed_no_obj_ptr: true
+  # multimask tracking settings
+  multimask_output_for_tracking: true
+  use_multimask_token_for_obj_ptr: true
+  multimask_min_pt_num: 0
+  multimask_max_pt_num: 1
+  use_mlp_for_obj_ptr_proj: true
+  # Compilation flag
+  compile_image_encoder: False

utility/sam/configs/sam2.1/sam2.1_hiera_s.yaml

@@ -0,0 +1,119 @@
+# @package _global_
+
+# Model
+model:
+  _target_: sam2.modeling.sam2_base.SAM2Base
+  image_encoder:
+    _target_: sam2.modeling.backbones.image_encoder.ImageEncoder
+    scalp: 1
+    trunk:
+      _target_: sam2.modeling.backbones.hieradet.Hiera
+      embed_dim: 96
+      num_heads: 1
+      stages: [1, 2, 11, 2]
+      global_att_blocks: [7, 10, 13]
+      window_pos_embed_bkg_spatial_size: [7, 7]
+    neck:
+      _target_: sam2.modeling.backbones.image_encoder.FpnNeck
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 256
+        normalize: true
+        scale: null
+        temperature: 10000
+      d_model: 256
+      backbone_channel_list: [768, 384, 192, 96]
+      fpn_top_down_levels: [2, 3]  # output level 0 and 1 directly use the backbone features
+      fpn_interp_model: nearest
+
+  memory_attention:
+    _target_: sam2.modeling.memory_attention.MemoryAttention
+    d_model: 256
+    pos_enc_at_input: true
+    layer:
+      _target_: sam2.modeling.memory_attention.MemoryAttentionLayer
+      activation: relu
+      dim_feedforward: 2048
+      dropout: 0.1
+      pos_enc_at_attn: false
+      self_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [64, 64]
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+      d_model: 256
+      pos_enc_at_cross_attn_keys: true
+      pos_enc_at_cross_attn_queries: false
+      cross_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [64, 64]
+        rope_k_repeat: True
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+        kv_in_dim: 64
+    num_layers: 4
+
+  memory_encoder:
+      _target_: sam2.modeling.memory_encoder.MemoryEncoder
+      out_dim: 64
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 64
+        normalize: true
+        scale: null
+        temperature: 10000
+      mask_downsampler:
+        _target_: sam2.modeling.memory_encoder.MaskDownSampler
+        kernel_size: 3
+        stride: 2
+        padding: 1
+      fuser:
+        _target_: sam2.modeling.memory_encoder.Fuser
+        layer:
+          _target_: sam2.modeling.memory_encoder.CXBlock
+          dim: 256
+          kernel_size: 7
+          padding: 3
+          layer_scale_init_value: 1e-6
+          use_dwconv: True  # depth-wise convs
+        num_layers: 2
+
+  num_maskmem: 7
+  image_size: 1024
+  # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask
+  sigmoid_scale_for_mem_enc: 20.0
+  sigmoid_bias_for_mem_enc: -10.0
+  use_mask_input_as_output_without_sam: true
+  # Memory
+  directly_add_no_mem_embed: true
+  no_obj_embed_spatial: true
+  # use high-resolution feature map in the SAM mask decoder
+  use_high_res_features_in_sam: true
+  # output 3 masks on the first click on initial conditioning frames
+  multimask_output_in_sam: true
+  # SAM heads
+  iou_prediction_use_sigmoid: True
+  # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder
+  use_obj_ptrs_in_encoder: true
+  add_tpos_enc_to_obj_ptrs: true
+  proj_tpos_enc_in_obj_ptrs: true
+  use_signed_tpos_enc_to_obj_ptrs: true
+  only_obj_ptrs_in_the_past_for_eval: true
+  # object occlusion prediction
+  pred_obj_scores: true
+  pred_obj_scores_mlp: true
+  fixed_no_obj_ptr: true
+  # multimask tracking settings
+  multimask_output_for_tracking: true
+  use_multimask_token_for_obj_ptr: true
+  multimask_min_pt_num: 0
+  multimask_max_pt_num: 1
+  use_mlp_for_obj_ptr_proj: true
+  # Compilation flag
+  compile_image_encoder: False

utility/sam/configs/sam2.1/sam2.1_hiera_t.yaml

@@ -0,0 +1,121 @@
+# @package _global_
+
+# Model
+model:
+  _target_: sam2.modeling.sam2_base.SAM2Base
+  image_encoder:
+    _target_: sam2.modeling.backbones.image_encoder.ImageEncoder
+    scalp: 1
+    trunk:
+      _target_: sam2.modeling.backbones.hieradet.Hiera
+      embed_dim: 96
+      num_heads: 1
+      stages: [1, 2, 7, 2]
+      global_att_blocks: [5, 7, 9]
+      window_pos_embed_bkg_spatial_size: [7, 7]
+    neck:
+      _target_: sam2.modeling.backbones.image_encoder.FpnNeck
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 256
+        normalize: true
+        scale: null
+        temperature: 10000
+      d_model: 256
+      backbone_channel_list: [768, 384, 192, 96]
+      fpn_top_down_levels: [2, 3]  # output level 0 and 1 directly use the backbone features
+      fpn_interp_model: nearest
+
+  memory_attention:
+    _target_: sam2.modeling.memory_attention.MemoryAttention
+    d_model: 256
+    pos_enc_at_input: true
+    layer:
+      _target_: sam2.modeling.memory_attention.MemoryAttentionLayer
+      activation: relu
+      dim_feedforward: 2048
+      dropout: 0.1
+      pos_enc_at_attn: false
+      self_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [64, 64]
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+      d_model: 256
+      pos_enc_at_cross_attn_keys: true
+      pos_enc_at_cross_attn_queries: false
+      cross_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [64, 64]
+        rope_k_repeat: True
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+        kv_in_dim: 64
+    num_layers: 4
+
+  memory_encoder:
+      _target_: sam2.modeling.memory_encoder.MemoryEncoder
+      out_dim: 64
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 64
+        normalize: true
+        scale: null
+        temperature: 10000
+      mask_downsampler:
+        _target_: sam2.modeling.memory_encoder.MaskDownSampler
+        kernel_size: 3
+        stride: 2
+        padding: 1
+      fuser:
+        _target_: sam2.modeling.memory_encoder.Fuser
+        layer:
+          _target_: sam2.modeling.memory_encoder.CXBlock
+          dim: 256
+          kernel_size: 7
+          padding: 3
+          layer_scale_init_value: 1e-6
+          use_dwconv: True  # depth-wise convs
+        num_layers: 2
+
+  num_maskmem: 7
+  image_size: 1024
+  # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask
+  # SAM decoder
+  sigmoid_scale_for_mem_enc: 20.0
+  sigmoid_bias_for_mem_enc: -10.0
+  use_mask_input_as_output_without_sam: true
+  # Memory
+  directly_add_no_mem_embed: true
+  no_obj_embed_spatial: true
+  # use high-resolution feature map in the SAM mask decoder
+  use_high_res_features_in_sam: true
+  # output 3 masks on the first click on initial conditioning frames
+  multimask_output_in_sam: true
+  # SAM heads
+  iou_prediction_use_sigmoid: True
+  # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder
+  use_obj_ptrs_in_encoder: true
+  add_tpos_enc_to_obj_ptrs: true
+  proj_tpos_enc_in_obj_ptrs: true
+  use_signed_tpos_enc_to_obj_ptrs: true
+  only_obj_ptrs_in_the_past_for_eval: true
+  # object occlusion prediction
+  pred_obj_scores: true
+  pred_obj_scores_mlp: true
+  fixed_no_obj_ptr: true
+  # multimask tracking settings
+  multimask_output_for_tracking: true
+  use_multimask_token_for_obj_ptr: true
+  multimask_min_pt_num: 0
+  multimask_max_pt_num: 1
+  use_mlp_for_obj_ptr_proj: true
+  # Compilation flag
+  # HieraT does not currently support compilation, should always be set to False
+  compile_image_encoder: False

utility/sam/configs/sam2.1_training/sam2.1_hiera_b+_MOSE_finetune.yaml

@@ -0,0 +1,339 @@
+# @package _global_
+
+scratch:
+  resolution: 1024
+  train_batch_size: 1
+  num_train_workers: 10
+  num_frames: 8
+  max_num_objects: 3
+  base_lr: 5.0e-6
+  vision_lr: 3.0e-06
+  phases_per_epoch: 1
+  num_epochs: 40
+
+dataset:
+  # PATHS to Dataset
+  img_folder: null # PATH to MOSE JPEGImages folder
+  gt_folder: null  # PATH to MOSE Annotations folder
+  file_list_txt: training/assets/MOSE_sample_train_list.txt # Optional PATH to filelist containing a subset of videos to be used for training
+  multiplier: 2
+
+# Video transforms
+vos:
+  train_transforms:
+    - _target_: training.dataset.transforms.ComposeAPI
+      transforms:
+        - _target_: training.dataset.transforms.RandomHorizontalFlip
+          consistent_transform: True
+        - _target_: training.dataset.transforms.RandomAffine
+          degrees: 25
+          shear: 20
+          image_interpolation: bilinear
+          consistent_transform: True
+        - _target_: training.dataset.transforms.RandomResizeAPI
+          sizes: ${scratch.resolution}
+          square: true
+          consistent_transform: True
+        - _target_: training.dataset.transforms.ColorJitter
+          consistent_transform: True
+          brightness: 0.1
+          contrast: 0.03
+          saturation: 0.03
+          hue: null
+        - _target_: training.dataset.transforms.RandomGrayscale
+          p: 0.05
+          consistent_transform: True
+        - _target_: training.dataset.transforms.ColorJitter
+          consistent_transform: False
+          brightness: 0.1
+          contrast: 0.05
+          saturation: 0.05
+          hue: null
+        - _target_: training.dataset.transforms.ToTensorAPI
+        - _target_: training.dataset.transforms.NormalizeAPI
+          mean: [0.485, 0.456, 0.406]
+          std: [0.229, 0.224, 0.225]
+
+trainer:
+  _target_: training.trainer.Trainer
+  mode: train_only
+  max_epochs: ${times:${scratch.num_epochs},${scratch.phases_per_epoch}}
+  accelerator: cuda
+  seed_value: 123
+
+  model:
+    _target_: training.model.sam2.SAM2Train
+    image_encoder:
+      _target_: sam2.modeling.backbones.image_encoder.ImageEncoder
+      scalp: 1
+      trunk:
+        _target_: sam2.modeling.backbones.hieradet.Hiera
+        embed_dim: 112
+        num_heads: 2
+        drop_path_rate: 0.1
+      neck:
+        _target_: sam2.modeling.backbones.image_encoder.FpnNeck
+        position_encoding:
+          _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+          num_pos_feats: 256
+          normalize: true
+          scale: null
+          temperature: 10000
+        d_model: 256
+        backbone_channel_list: [896, 448, 224, 112]
+        fpn_top_down_levels: [2, 3]  # output level 0 and 1 directly use the backbone features
+        fpn_interp_model: nearest
+
+    memory_attention:
+      _target_: sam2.modeling.memory_attention.MemoryAttention
+      d_model: 256
+      pos_enc_at_input: true
+      layer:
+        _target_: sam2.modeling.memory_attention.MemoryAttentionLayer
+        activation: relu
+        dim_feedforward: 2048
+        dropout: 0.1
+        pos_enc_at_attn: false
+        self_attention:
+          _target_: sam2.modeling.sam.transformer.RoPEAttention
+          rope_theta: 10000.0
+          feat_sizes: [64, 64]
+          embedding_dim: 256
+          num_heads: 1
+          downsample_rate: 1
+          dropout: 0.1
+        d_model: 256
+        pos_enc_at_cross_attn_keys: true
+        pos_enc_at_cross_attn_queries: false
+        cross_attention:
+          _target_: sam2.modeling.sam.transformer.RoPEAttention
+          rope_theta: 10000.0
+          feat_sizes: [64, 64]
+          rope_k_repeat: True
+          embedding_dim: 256
+          num_heads: 1
+          downsample_rate: 1
+          dropout: 0.1
+          kv_in_dim: 64
+      num_layers: 4
+
+    memory_encoder:
+        _target_: sam2.modeling.memory_encoder.MemoryEncoder
+        out_dim: 64
+        position_encoding:
+          _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+          num_pos_feats: 64
+          normalize: true
+          scale: null
+          temperature: 10000
+        mask_downsampler:
+          _target_: sam2.modeling.memory_encoder.MaskDownSampler
+          kernel_size: 3
+          stride: 2
+          padding: 1
+        fuser:
+          _target_: sam2.modeling.memory_encoder.Fuser
+          layer:
+            _target_: sam2.modeling.memory_encoder.CXBlock
+            dim: 256
+            kernel_size: 7
+            padding: 3
+            layer_scale_init_value: 1e-6
+            use_dwconv: True  # depth-wise convs
+          num_layers: 2
+
+    num_maskmem: 7
+    image_size: ${scratch.resolution}
+    # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask
+    sigmoid_scale_for_mem_enc: 20.0
+    sigmoid_bias_for_mem_enc: -10.0
+    use_mask_input_as_output_without_sam: true
+    # Memory
+    directly_add_no_mem_embed: true
+    no_obj_embed_spatial: true
+    # use high-resolution feature map in the SAM mask decoder
+    use_high_res_features_in_sam: true
+    # output 3 masks on the first click on initial conditioning frames
+    multimask_output_in_sam: true
+    # SAM heads
+    iou_prediction_use_sigmoid: True
+    # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder
+    use_obj_ptrs_in_encoder: true
+    add_tpos_enc_to_obj_ptrs: true
+    proj_tpos_enc_in_obj_ptrs: true
+    use_signed_tpos_enc_to_obj_ptrs: true
+    only_obj_ptrs_in_the_past_for_eval: true
+    # object occlusion prediction
+    pred_obj_scores: true
+    pred_obj_scores_mlp: true
+    fixed_no_obj_ptr: true
+    # multimask tracking settings
+    multimask_output_for_tracking: true
+    use_multimask_token_for_obj_ptr: true
+    multimask_min_pt_num: 0
+    multimask_max_pt_num: 1
+    use_mlp_for_obj_ptr_proj: true
+    # Compilation flag
+    # compile_image_encoder: False
+
+    ####### Training specific params #######
+    # box/point input and corrections
+    prob_to_use_pt_input_for_train: 0.5
+    prob_to_use_pt_input_for_eval: 0.0
+    prob_to_use_box_input_for_train: 0.5  # 0.5*0.5 = 0.25 prob to use box instead of points
+    prob_to_use_box_input_for_eval: 0.0
+    prob_to_sample_from_gt_for_train: 0.1  # with a small prob, sampling correction points from GT mask instead of prediction errors
+    num_frames_to_correct_for_train: 2  # iteratively sample on random 1~2 frames (always include the first frame)
+    num_frames_to_correct_for_eval: 1  # only iteratively sample on first frame
+    rand_frames_to_correct_for_train: True  # random #init-cond-frame ~ 2
+    add_all_frames_to_correct_as_cond: True  # when a frame receives a correction click, it becomes a conditioning frame (even if it's not initially a conditioning frame)
+    # maximum 2 initial conditioning frames
+    num_init_cond_frames_for_train: 2
+    rand_init_cond_frames_for_train: True  # random 1~2
+    num_correction_pt_per_frame: 7
+    use_act_ckpt_iterative_pt_sampling: false
+    
+
+    
+    num_init_cond_frames_for_eval: 1  # only mask on the first frame
+    forward_backbone_per_frame_for_eval: True
+    
+
+  data:
+    train:
+      _target_: training.dataset.sam2_datasets.TorchTrainMixedDataset
+      phases_per_epoch: ${scratch.phases_per_epoch}
+      batch_sizes:
+        - ${scratch.train_batch_size}
+
+      datasets:
+        - _target_: training.dataset.utils.RepeatFactorWrapper
+          dataset:
+            _target_: training.dataset.utils.ConcatDataset
+            datasets:
+            - _target_: training.dataset.vos_dataset.VOSDataset
+              transforms: ${vos.train_transforms}
+              training: true
+              video_dataset:
+                _target_: training.dataset.vos_raw_dataset.PNGRawDataset
+                img_folder: ${dataset.img_folder}
+                gt_folder: ${dataset.gt_folder}
+                file_list_txt: ${dataset.file_list_txt}
+              sampler:
+                _target_: training.dataset.vos_sampler.RandomUniformSampler
+                num_frames: ${scratch.num_frames}
+                max_num_objects: ${scratch.max_num_objects}
+              multiplier: ${dataset.multiplier}
+      shuffle: True
+      num_workers: ${scratch.num_train_workers}
+      pin_memory: True
+      drop_last: True
+      collate_fn:
+        _target_: training.utils.data_utils.collate_fn
+        _partial_: true
+        dict_key: all
+
+  optim:
+    amp:
+      enabled: True
+      amp_dtype: bfloat16
+
+    optimizer:
+      _target_: torch.optim.AdamW
+
+    gradient_clip:
+      _target_: training.optimizer.GradientClipper
+      max_norm: 0.1
+      norm_type: 2
+
+    param_group_modifiers:
+      - _target_: training.optimizer.layer_decay_param_modifier
+        _partial_: True
+        layer_decay_value: 0.9
+        apply_to: 'image_encoder.trunk'
+        overrides:
+          - pattern: '*pos_embed*'
+            value: 1.0
+
+    options:
+      lr:
+        - scheduler:
+            _target_: fvcore.common.param_scheduler.CosineParamScheduler
+            start_value: ${scratch.base_lr}
+            end_value: ${divide:${scratch.base_lr},10}
+        - scheduler:
+            _target_: fvcore.common.param_scheduler.CosineParamScheduler
+            start_value: ${scratch.vision_lr}
+            end_value: ${divide:${scratch.vision_lr},10}
+          param_names:
+            - 'image_encoder.*'
+      weight_decay:
+        - scheduler:
+            _target_: fvcore.common.param_scheduler.ConstantParamScheduler
+            value: 0.1
+        - scheduler:
+            _target_: fvcore.common.param_scheduler.ConstantParamScheduler
+            value: 0.0
+          param_names:
+            - '*bias*'
+          module_cls_names: ['torch.nn.LayerNorm']
+
+  loss:
+    all:
+      _target_: training.loss_fns.MultiStepMultiMasksAndIous
+      weight_dict:
+        loss_mask: 20
+        loss_dice: 1
+        loss_iou: 1
+        loss_class: 1
+      supervise_all_iou: true
+      iou_use_l1_loss: true
+      pred_obj_scores: true
+      focal_gamma_obj_score: 0.0
+      focal_alpha_obj_score: -1.0
+
+  distributed:
+    backend: nccl
+    find_unused_parameters: True
+
+  logging:
+    tensorboard_writer:
+      _target_: training.utils.logger.make_tensorboard_logger
+      log_dir:  ${launcher.experiment_log_dir}/tensorboard
+      flush_secs: 120
+      should_log: True
+    log_dir: ${launcher.experiment_log_dir}/logs
+    log_freq: 10
+
+  # initialize from a SAM 2 checkpoint
+  checkpoint:
+    save_dir: ${launcher.experiment_log_dir}/checkpoints
+    save_freq: 0 # 0 only last checkpoint is saved.
+    model_weight_initializer:
+      _partial_: True
+      _target_: training.utils.checkpoint_utils.load_state_dict_into_model
+      strict: True
+      ignore_unexpected_keys: null
+      ignore_missing_keys: null
+
+      state_dict:
+        _target_: training.utils.checkpoint_utils.load_checkpoint_and_apply_kernels
+        checkpoint_path: ./checkpoints/sam2.1_hiera_base_plus.pt # PATH to SAM 2.1 checkpoint
+        ckpt_state_dict_keys: ['model']
+
+launcher:
+  num_nodes: 1
+  gpus_per_node: 8
+  experiment_log_dir: null # Path to log directory, defaults to ./sam2_logs/${config_name}
+
+# SLURM args if running on a cluster
+submitit:
+  partition: null
+  account: null
+  qos: null
+  cpus_per_task: 10
+  use_cluster: false
+  timeout_hour: 24
+  name: null
+  port_range: [10000, 65000]
+

utility/sam/configs/sam2/sam2_hiera_b+.yaml

@@ -0,0 +1,113 @@
+# @package _global_
+
+# Model
+model:
+  _target_: sam2.modeling.sam2_base.SAM2Base
+  image_encoder:
+    _target_: sam2.modeling.backbones.image_encoder.ImageEncoder
+    scalp: 1
+    trunk:
+      _target_: sam2.modeling.backbones.hieradet.Hiera
+      embed_dim: 112
+      num_heads: 2
+    neck:
+      _target_: sam2.modeling.backbones.image_encoder.FpnNeck
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 256
+        normalize: true
+        scale: null
+        temperature: 10000
+      d_model: 256
+      backbone_channel_list: [896, 448, 224, 112]
+      fpn_top_down_levels: [2, 3]  # output level 0 and 1 directly use the backbone features
+      fpn_interp_model: nearest
+
+  memory_attention:
+    _target_: sam2.modeling.memory_attention.MemoryAttention
+    d_model: 256
+    pos_enc_at_input: true
+    layer:
+      _target_: sam2.modeling.memory_attention.MemoryAttentionLayer
+      activation: relu
+      dim_feedforward: 2048
+      dropout: 0.1
+      pos_enc_at_attn: false
+      self_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [64, 64]
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+      d_model: 256
+      pos_enc_at_cross_attn_keys: true
+      pos_enc_at_cross_attn_queries: false
+      cross_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [64, 64]
+        rope_k_repeat: True
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+        kv_in_dim: 64
+    num_layers: 4
+
+  memory_encoder:
+      _target_: sam2.modeling.memory_encoder.MemoryEncoder
+      out_dim: 64
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 64
+        normalize: true
+        scale: null
+        temperature: 10000
+      mask_downsampler:
+        _target_: sam2.modeling.memory_encoder.MaskDownSampler
+        kernel_size: 3
+        stride: 2
+        padding: 1
+      fuser:
+        _target_: sam2.modeling.memory_encoder.Fuser
+        layer:
+          _target_: sam2.modeling.memory_encoder.CXBlock
+          dim: 256
+          kernel_size: 7
+          padding: 3
+          layer_scale_init_value: 1e-6
+          use_dwconv: True  # depth-wise convs
+        num_layers: 2
+
+  num_maskmem: 7
+  image_size: 1024
+  # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask
+  sigmoid_scale_for_mem_enc: 20.0
+  sigmoid_bias_for_mem_enc: -10.0
+  use_mask_input_as_output_without_sam: true
+  # Memory
+  directly_add_no_mem_embed: true
+  # use high-resolution feature map in the SAM mask decoder
+  use_high_res_features_in_sam: true
+  # output 3 masks on the first click on initial conditioning frames
+  multimask_output_in_sam: true
+  # SAM heads
+  iou_prediction_use_sigmoid: True
+  # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder
+  use_obj_ptrs_in_encoder: true
+  add_tpos_enc_to_obj_ptrs: false
+  only_obj_ptrs_in_the_past_for_eval: true
+  # object occlusion prediction
+  pred_obj_scores: true
+  pred_obj_scores_mlp: true
+  fixed_no_obj_ptr: true
+  # multimask tracking settings
+  multimask_output_for_tracking: true
+  use_multimask_token_for_obj_ptr: true
+  multimask_min_pt_num: 0
+  multimask_max_pt_num: 1
+  use_mlp_for_obj_ptr_proj: true
+  # Compilation flag
+  compile_image_encoder: False

utility/sam/configs/sam2/sam2_hiera_l.yaml

@@ -0,0 +1,117 @@
+# @package _global_
+
+# Model
+model:
+  _target_: sam2.modeling.sam2_base.SAM2Base
+  image_encoder:
+    _target_: sam2.modeling.backbones.image_encoder.ImageEncoder
+    scalp: 1
+    trunk:
+      _target_: sam2.modeling.backbones.hieradet.Hiera
+      embed_dim: 144
+      num_heads: 2
+      stages: [2, 6, 36, 4]
+      global_att_blocks: [23, 33, 43]
+      window_pos_embed_bkg_spatial_size: [7, 7]
+      window_spec: [8, 4, 16, 8]
+    neck:
+      _target_: sam2.modeling.backbones.image_encoder.FpnNeck
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 256
+        normalize: true
+        scale: null
+        temperature: 10000
+      d_model: 256
+      backbone_channel_list: [1152, 576, 288, 144]
+      fpn_top_down_levels: [2, 3]  # output level 0 and 1 directly use the backbone features
+      fpn_interp_model: nearest
+
+  memory_attention:
+    _target_: sam2.modeling.memory_attention.MemoryAttention
+    d_model: 256
+    pos_enc_at_input: true
+    layer:
+      _target_: sam2.modeling.memory_attention.MemoryAttentionLayer
+      activation: relu
+      dim_feedforward: 2048
+      dropout: 0.1
+      pos_enc_at_attn: false
+      self_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [64, 64]
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+      d_model: 256
+      pos_enc_at_cross_attn_keys: true
+      pos_enc_at_cross_attn_queries: false
+      cross_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [64, 64]
+        rope_k_repeat: True
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+        kv_in_dim: 64
+    num_layers: 4
+
+  memory_encoder:
+      _target_: sam2.modeling.memory_encoder.MemoryEncoder
+      out_dim: 64
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 64
+        normalize: true
+        scale: null
+        temperature: 10000
+      mask_downsampler:
+        _target_: sam2.modeling.memory_encoder.MaskDownSampler
+        kernel_size: 3
+        stride: 2
+        padding: 1
+      fuser:
+        _target_: sam2.modeling.memory_encoder.Fuser
+        layer:
+          _target_: sam2.modeling.memory_encoder.CXBlock
+          dim: 256
+          kernel_size: 7
+          padding: 3
+          layer_scale_init_value: 1e-6
+          use_dwconv: True  # depth-wise convs
+        num_layers: 2
+
+  num_maskmem: 7
+  image_size: 1024
+  # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask
+  sigmoid_scale_for_mem_enc: 20.0
+  sigmoid_bias_for_mem_enc: -10.0
+  use_mask_input_as_output_without_sam: true
+  # Memory
+  directly_add_no_mem_embed: true
+  # use high-resolution feature map in the SAM mask decoder
+  use_high_res_features_in_sam: true
+  # output 3 masks on the first click on initial conditioning frames
+  multimask_output_in_sam: true
+  # SAM heads
+  iou_prediction_use_sigmoid: True
+  # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder
+  use_obj_ptrs_in_encoder: true
+  add_tpos_enc_to_obj_ptrs: false
+  only_obj_ptrs_in_the_past_for_eval: true
+  # object occlusion prediction
+  pred_obj_scores: true
+  pred_obj_scores_mlp: true
+  fixed_no_obj_ptr: true
+  # multimask tracking settings
+  multimask_output_for_tracking: true
+  use_multimask_token_for_obj_ptr: true
+  multimask_min_pt_num: 0
+  multimask_max_pt_num: 1
+  use_mlp_for_obj_ptr_proj: true
+  # Compilation flag
+  compile_image_encoder: False

utility/sam/configs/sam2/sam2_hiera_s.yaml

@@ -0,0 +1,116 @@
+# @package _global_
+
+# Model
+model:
+  _target_: sam2.modeling.sam2_base.SAM2Base
+  image_encoder:
+    _target_: sam2.modeling.backbones.image_encoder.ImageEncoder
+    scalp: 1
+    trunk:
+      _target_: sam2.modeling.backbones.hieradet.Hiera
+      embed_dim: 96
+      num_heads: 1
+      stages: [1, 2, 11, 2]
+      global_att_blocks: [7, 10, 13]
+      window_pos_embed_bkg_spatial_size: [7, 7]
+    neck:
+      _target_: sam2.modeling.backbones.image_encoder.FpnNeck
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 256
+        normalize: true
+        scale: null
+        temperature: 10000
+      d_model: 256
+      backbone_channel_list: [768, 384, 192, 96]
+      fpn_top_down_levels: [2, 3]  # output level 0 and 1 directly use the backbone features
+      fpn_interp_model: nearest
+
+  memory_attention:
+    _target_: sam2.modeling.memory_attention.MemoryAttention
+    d_model: 256
+    pos_enc_at_input: true
+    layer:
+      _target_: sam2.modeling.memory_attention.MemoryAttentionLayer
+      activation: relu
+      dim_feedforward: 2048
+      dropout: 0.1
+      pos_enc_at_attn: false
+      self_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [64, 64]
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+      d_model: 256
+      pos_enc_at_cross_attn_keys: true
+      pos_enc_at_cross_attn_queries: false
+      cross_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [64, 64]
+        rope_k_repeat: True
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+        kv_in_dim: 64
+    num_layers: 4
+
+  memory_encoder:
+      _target_: sam2.modeling.memory_encoder.MemoryEncoder
+      out_dim: 64
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 64
+        normalize: true
+        scale: null
+        temperature: 10000
+      mask_downsampler:
+        _target_: sam2.modeling.memory_encoder.MaskDownSampler
+        kernel_size: 3
+        stride: 2
+        padding: 1
+      fuser:
+        _target_: sam2.modeling.memory_encoder.Fuser
+        layer:
+          _target_: sam2.modeling.memory_encoder.CXBlock
+          dim: 256
+          kernel_size: 7
+          padding: 3
+          layer_scale_init_value: 1e-6
+          use_dwconv: True  # depth-wise convs
+        num_layers: 2
+
+  num_maskmem: 7
+  image_size: 1024
+  # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask
+  sigmoid_scale_for_mem_enc: 20.0
+  sigmoid_bias_for_mem_enc: -10.0
+  use_mask_input_as_output_without_sam: true
+  # Memory
+  directly_add_no_mem_embed: true
+  # use high-resolution feature map in the SAM mask decoder
+  use_high_res_features_in_sam: true
+  # output 3 masks on the first click on initial conditioning frames
+  multimask_output_in_sam: true
+  # SAM heads
+  iou_prediction_use_sigmoid: True
+  # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder
+  use_obj_ptrs_in_encoder: true
+  add_tpos_enc_to_obj_ptrs: false
+  only_obj_ptrs_in_the_past_for_eval: true
+  # object occlusion prediction
+  pred_obj_scores: true
+  pred_obj_scores_mlp: true
+  fixed_no_obj_ptr: true
+  # multimask tracking settings
+  multimask_output_for_tracking: true
+  use_multimask_token_for_obj_ptr: true
+  multimask_min_pt_num: 0
+  multimask_max_pt_num: 1
+  use_mlp_for_obj_ptr_proj: true
+  # Compilation flag
+  compile_image_encoder: False

utility/sam/configs/sam2/sam2_hiera_t.yaml

@@ -0,0 +1,118 @@
+# @package _global_
+
+# Model
+model:
+  _target_: sam2.modeling.sam2_base.SAM2Base
+  image_encoder:
+    _target_: sam2.modeling.backbones.image_encoder.ImageEncoder
+    scalp: 1
+    trunk:
+      _target_: sam2.modeling.backbones.hieradet.Hiera
+      embed_dim: 96
+      num_heads: 1
+      stages: [1, 2, 7, 2]
+      global_att_blocks: [5, 7, 9]
+      window_pos_embed_bkg_spatial_size: [7, 7]
+    neck:
+      _target_: sam2.modeling.backbones.image_encoder.FpnNeck
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 256
+        normalize: true
+        scale: null
+        temperature: 10000
+      d_model: 256
+      backbone_channel_list: [768, 384, 192, 96]
+      fpn_top_down_levels: [2, 3]  # output level 0 and 1 directly use the backbone features
+      fpn_interp_model: nearest
+
+  memory_attention:
+    _target_: sam2.modeling.memory_attention.MemoryAttention
+    d_model: 256
+    pos_enc_at_input: true
+    layer:
+      _target_: sam2.modeling.memory_attention.MemoryAttentionLayer
+      activation: relu
+      dim_feedforward: 2048
+      dropout: 0.1
+      pos_enc_at_attn: false
+      self_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [64, 64]
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+      d_model: 256
+      pos_enc_at_cross_attn_keys: true
+      pos_enc_at_cross_attn_queries: false
+      cross_attention:
+        _target_: sam2.modeling.sam.transformer.RoPEAttention
+        rope_theta: 10000.0
+        feat_sizes: [64, 64]
+        rope_k_repeat: True
+        embedding_dim: 256
+        num_heads: 1
+        downsample_rate: 1
+        dropout: 0.1
+        kv_in_dim: 64
+    num_layers: 4
+
+  memory_encoder:
+      _target_: sam2.modeling.memory_encoder.MemoryEncoder
+      out_dim: 64
+      position_encoding:
+        _target_: sam2.modeling.position_encoding.PositionEmbeddingSine
+        num_pos_feats: 64
+        normalize: true
+        scale: null
+        temperature: 10000
+      mask_downsampler:
+        _target_: sam2.modeling.memory_encoder.MaskDownSampler
+        kernel_size: 3
+        stride: 2
+        padding: 1
+      fuser:
+        _target_: sam2.modeling.memory_encoder.Fuser
+        layer:
+          _target_: sam2.modeling.memory_encoder.CXBlock
+          dim: 256
+          kernel_size: 7
+          padding: 3
+          layer_scale_init_value: 1e-6
+          use_dwconv: True  # depth-wise convs
+        num_layers: 2
+
+  num_maskmem: 7
+  image_size: 1024
+  # apply scaled sigmoid on mask logits for memory encoder, and directly feed input mask as output mask
+  # SAM decoder
+  sigmoid_scale_for_mem_enc: 20.0
+  sigmoid_bias_for_mem_enc: -10.0
+  use_mask_input_as_output_without_sam: true
+  # Memory
+  directly_add_no_mem_embed: true
+  # use high-resolution feature map in the SAM mask decoder
+  use_high_res_features_in_sam: true
+  # output 3 masks on the first click on initial conditioning frames
+  multimask_output_in_sam: true
+  # SAM heads
+  iou_prediction_use_sigmoid: True
+  # cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder
+  use_obj_ptrs_in_encoder: true
+  add_tpos_enc_to_obj_ptrs: false
+  only_obj_ptrs_in_the_past_for_eval: true
+  # object occlusion prediction
+  pred_obj_scores: true
+  pred_obj_scores_mlp: true
+  fixed_no_obj_ptr: true
+  # multimask tracking settings
+  multimask_output_for_tracking: true
+  use_multimask_token_for_obj_ptr: true
+  multimask_min_pt_num: 0
+  multimask_max_pt_num: 1
+  use_mlp_for_obj_ptr_proj: true
+  # Compilation flag
+  # HieraT does not currently support compilation, should always be set to False
+  compile_image_encoder: False

utility/sam/modeling/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.

utility/sam/modeling/backbones/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.

utility/sam/modeling/backbones/hieradet.py

@@ -0,0 +1,318 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+from functools import partial
+from typing import List, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+#from iopath.common.file_io import g_pathmgr
+
+from sam2.modeling.backbones.utils import (
+    PatchEmbed,
+    window_partition,
+    window_unpartition,
+)
+
+from sam2.modeling.sam2_utils import DropPath, MLP
+
+
+def do_pool(x: torch.Tensor, pool: nn.Module, norm: nn.Module = None) -> torch.Tensor:
+    if pool is None:
+        return x
+    # (B, H, W, C) -> (B, C, H, W)
+    x = x.permute(0, 3, 1, 2)
+    x = pool(x)
+    # (B, C, H', W') -> (B, H', W', C)
+    x = x.permute(0, 2, 3, 1)
+    if norm:
+        x = norm(x)
+
+    return x
+
+
+class MultiScaleAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        dim_out: int,
+        num_heads: int,
+        q_pool: nn.Module = None,
+    ):
+        super().__init__()
+
+        self.dim = dim
+        self.dim_out = dim_out
+        self.num_heads = num_heads
+        self.q_pool = q_pool
+        self.qkv = nn.Linear(dim, dim_out * 3)
+        self.proj = nn.Linear(dim_out, dim_out)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, H, W, _ = x.shape
+        # qkv with shape (B, H * W, 3, nHead, C)
+        qkv = self.qkv(x).reshape(B, H * W, 3, self.num_heads, -1)
+        # q, k, v with shape (B, H * W, nheads, C)
+        q, k, v = torch.unbind(qkv, 2)
+
+        # Q pooling (for downsample at stage changes)
+        if self.q_pool:
+            q = do_pool(q.reshape(B, H, W, -1), self.q_pool)
+            H, W = q.shape[1:3]  # downsampled shape
+            q = q.reshape(B, H * W, self.num_heads, -1)
+
+        # Torch's SDPA expects [B, nheads, H*W, C] so we transpose
+        x = F.scaled_dot_product_attention(
+            q.transpose(1, 2),
+            k.transpose(1, 2),
+            v.transpose(1, 2),
+        )
+        # Transpose back
+        x = x.transpose(1, 2)
+        x = x.reshape(B, H, W, -1)
+
+        x = self.proj(x)
+
+        return x
+
+
+class MultiScaleBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        dim_out: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        drop_path: float = 0.0,
+        norm_layer: Union[nn.Module, str] = "LayerNorm",
+        q_stride: Tuple[int, int] = None,
+        act_layer: nn.Module = nn.GELU,
+        window_size: int = 0,
+    ):
+        super().__init__()
+
+        if isinstance(norm_layer, str):
+            norm_layer = partial(getattr(nn, norm_layer), eps=1e-6)
+
+        self.dim = dim
+        self.dim_out = dim_out
+        self.norm1 = norm_layer(dim)
+
+        self.window_size = window_size
+
+        self.pool, self.q_stride = None, q_stride
+        if self.q_stride:
+            self.pool = nn.MaxPool2d(
+                kernel_size=q_stride, stride=q_stride, ceil_mode=False
+            )
+
+        self.attn = MultiScaleAttention(
+            dim,
+            dim_out,
+            num_heads=num_heads,
+            q_pool=self.pool,
+        )
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        self.norm2 = norm_layer(dim_out)
+        self.mlp = MLP(
+            dim_out,
+            int(dim_out * mlp_ratio),
+            dim_out,
+            num_layers=2,
+            activation=act_layer,
+        )
+
+        if dim != dim_out:
+            self.proj = nn.Linear(dim, dim_out)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        shortcut = x  # B, H, W, C
+        x = self.norm1(x)
+
+        # Skip connection
+        if self.dim != self.dim_out:
+            shortcut = do_pool(self.proj(x), self.pool)
+
+        # Window partition
+        window_size = self.window_size
+        if window_size > 0:
+            H, W = x.shape[1], x.shape[2]
+            x, pad_hw = window_partition(x, window_size)
+
+        # Window Attention + Q Pooling (if stage change)
+        x = self.attn(x)
+        if self.q_stride:
+            # Shapes have changed due to Q pooling
+            window_size = self.window_size // self.q_stride[0]
+            H, W = shortcut.shape[1:3]
+
+            pad_h = (window_size - H % window_size) % window_size
+            pad_w = (window_size - W % window_size) % window_size
+            pad_hw = (H + pad_h, W + pad_w)
+
+        # Reverse window partition
+        if self.window_size > 0:
+            x = window_unpartition(x, window_size, pad_hw, (H, W))
+
+        x = shortcut + self.drop_path(x)
+        # MLP
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+
+class Hiera(nn.Module):
+    """
+    Reference: https://arxiv.org/abs/2306.00989
+    """
+
+    def __init__(
+        self,
+        embed_dim: int = 96,  # initial embed dim
+        num_heads: int = 1,  # initial number of heads
+        drop_path_rate: float = 0.0,  # stochastic depth
+        q_pool: int = 3,  # number of q_pool stages
+        q_stride: Tuple[int, int] = (2, 2),  # downsample stride bet. stages
+        stages: Tuple[int, ...] = (2, 3, 16, 3),  # blocks per stage
+        dim_mul: float = 2.0,  # dim_mul factor at stage shift
+        head_mul: float = 2.0,  # head_mul factor at stage shift
+        window_pos_embed_bkg_spatial_size: Tuple[int, int] = (14, 14),
+        # window size per stage, when not using global att.
+        window_spec: Tuple[int, ...] = (
+            8,
+            4,
+            14,
+            7,
+        ),
+        # global attn in these blocks
+        global_att_blocks: Tuple[int, ...] = (
+            12,
+            16,
+            20,
+        ),
+        weights_path=None,
+        return_interm_layers=True,  # return feats from every stage
+    ):
+        super().__init__()
+
+        assert len(stages) == len(window_spec)
+        self.window_spec = window_spec
+
+        depth = sum(stages)
+        self.q_stride = q_stride
+        self.stage_ends = [sum(stages[:i]) - 1 for i in range(1, len(stages) + 1)]
+        assert 0 <= q_pool <= len(self.stage_ends[:-1])
+        self.q_pool_blocks = [x + 1 for x in self.stage_ends[:-1]][:q_pool]
+        self.return_interm_layers = return_interm_layers
+
+        self.patch_embed = PatchEmbed(
+            embed_dim=embed_dim,
+        )
+        # Which blocks have global att?
+        self.global_att_blocks = global_att_blocks
+
+        # Windowed positional embedding (https://arxiv.org/abs/2311.05613)
+        self.window_pos_embed_bkg_spatial_size = window_pos_embed_bkg_spatial_size
+        self.pos_embed = nn.Parameter(
+            torch.zeros(1, embed_dim, *self.window_pos_embed_bkg_spatial_size)
+        )
+        self.pos_embed_window = nn.Parameter(
+            torch.zeros(1, embed_dim, self.window_spec[0], self.window_spec[0])
+        )
+
+        dpr = [
+            x.item() for x in torch.linspace(0, drop_path_rate, depth)
+        ]  # stochastic depth decay rule
+
+        cur_stage = 1
+        self.blocks = nn.ModuleList()
+
+        for i in range(depth):
+            dim_out = embed_dim
+            # lags by a block, so first block of
+            # next stage uses an initial window size
+            # of previous stage and final window size of current stage
+            window_size = self.window_spec[cur_stage - 1]
+
+            if self.global_att_blocks is not None:
+                window_size = 0 if i in self.global_att_blocks else window_size
+
+            if i - 1 in self.stage_ends:
+                dim_out = int(embed_dim * dim_mul)
+                num_heads = int(num_heads * head_mul)
+                cur_stage += 1
+
+            block = MultiScaleBlock(
+                dim=embed_dim,
+                dim_out=dim_out,
+                num_heads=num_heads,
+                drop_path=dpr[i],
+                q_stride=self.q_stride if i in self.q_pool_blocks else None,
+                window_size=window_size,
+            )
+
+            embed_dim = dim_out
+            self.blocks.append(block)
+
+        self.channel_list = (
+            [self.blocks[i].dim_out for i in self.stage_ends[::-1]]
+            if return_interm_layers
+            else [self.blocks[-1].dim_out]
+        )
+
+        if weights_path is not None:
+            #with g_pathmgr.open(weights_path, "rb") as f:
+            with open(weights_path, "rb") as f:
+                chkpt = torch.load(f, map_location="cpu")
+            self.load_state_dict(chkpt, strict=False)
+
+    def _get_pos_embed(self, hw: Tuple[int, int]) -> torch.Tensor:
+        h, w = hw
+        window_embed = self.pos_embed_window
+        pos_embed = F.interpolate(self.pos_embed, size=(h, w), mode="bicubic")
+        pos_embed = pos_embed + window_embed.tile(
+            [x // y for x, y in zip(pos_embed.shape, window_embed.shape)]
+        )
+        pos_embed = pos_embed.permute(0, 2, 3, 1)
+        return pos_embed
+
+    def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
+        x = self.patch_embed(x)
+        # x: (B, H, W, C)
+
+        # Add pos embed
+        x = x + self._get_pos_embed(x.shape[1:3])
+
+        outputs = []
+        for i, blk in enumerate(self.blocks):
+            x = blk(x)
+            if (i == self.stage_ends[-1]) or (
+                i in self.stage_ends and self.return_interm_layers
+            ):
+                feats = x.permute(0, 3, 1, 2)
+                outputs.append(feats)
+
+        return outputs
+
+    def get_layer_id(self, layer_name):
+        # https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py#L33
+        num_layers = self.get_num_layers()
+
+        if layer_name.find("rel_pos") != -1:
+            return num_layers + 1
+        elif layer_name.find("pos_embed") != -1:
+            return 0
+        elif layer_name.find("patch_embed") != -1:
+            return 0
+        elif layer_name.find("blocks") != -1:
+            return int(layer_name.split("blocks")[1].split(".")[1]) + 1
+        else:
+            return num_layers + 1
+
+    def get_num_layers(self) -> int:
+        return len(self.blocks)

utility/sam/modeling/backbones/image_encoder.py

@@ -0,0 +1,196 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import List, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from sam2.modeling.sam2_utils import LayerNorm2d
+
+class ImageEncoder(nn.Module):
+    def __init__(
+        self,
+        trunk: nn.Module,
+        neck: nn.Module,
+        scalp: int = 0,
+    ):
+        super().__init__()
+        self.trunk = trunk
+        self.neck = neck
+        self.scalp = scalp
+        assert (
+            self.trunk.channel_list == self.neck.backbone_channel_list
+        ), f"Channel dims of trunk and neck do not match. Trunk: {self.trunk.channel_list}, neck: {self.neck.backbone_channel_list}"
+
+    def forward(self, sample: torch.Tensor):
+        # Forward through backbone
+        features, pos = self.neck(self.trunk(sample))
+        if self.scalp > 0:
+            # Discard the lowest resolution features
+            features, pos = features[: -self.scalp], pos[: -self.scalp]
+
+        src = features[-1]
+        output = {
+            "vision_features": src,
+            "vision_pos_enc": pos,
+            "backbone_fpn": features,
+        }
+        return output
+
+
+class FpnNeck(nn.Module):
+    """
+    A modified variant of Feature Pyramid Network (FPN) neck
+    (we remove output conv and also do bicubic interpolation similar to ViT
+    pos embed interpolation)
+    """
+
+    def __init__(
+        self,
+        position_encoding: nn.Module,
+        d_model: int,
+        backbone_channel_list: List[int],
+        kernel_size: int = 1,
+        stride: int = 1,
+        padding: int = 0,
+        fpn_interp_model: str = "bilinear",
+        fuse_type: str = "sum",
+        fpn_top_down_levels: Optional[List[int]] = None,
+    ):
+        """Initialize the neck
+        :param trunk: the backbone
+        :param position_encoding: the positional encoding to use
+        :param d_model: the dimension of the model
+        :param neck_norm: the normalization to use
+        """
+        super().__init__()
+        self.position_encoding = position_encoding
+        self.convs = nn.ModuleList()
+        self.backbone_channel_list = backbone_channel_list
+        self.d_model = d_model
+        for dim in backbone_channel_list:
+            current = nn.Sequential()
+            current.add_module(
+                "conv",
+                nn.Conv2d(
+                    in_channels=dim,
+                    out_channels=d_model,
+                    kernel_size=kernel_size,
+                    stride=stride,
+                    padding=padding,
+                ),
+            )
+
+            self.convs.append(current)
+        self.fpn_interp_model = fpn_interp_model
+        assert fuse_type in ["sum", "avg"]
+        self.fuse_type = fuse_type
+
+        # levels to have top-down features in its outputs
+        # e.g. if fpn_top_down_levels is [2, 3], then only outputs of level 2 and 3
+        # have top-down propagation, while outputs of level 0 and level 1 have only
+        # lateral features from the same backbone level.
+        if fpn_top_down_levels is None:
+            # default is to have top-down features on all levels
+            fpn_top_down_levels = range(len(self.convs))
+        self.fpn_top_down_levels = list(fpn_top_down_levels)
+
+    def forward(self, xs: List[torch.Tensor]):
+
+        out = [None] * len(self.convs)
+        pos = [None] * len(self.convs)
+        assert len(xs) == len(self.convs)
+        # fpn forward pass
+        # see https://github.com/facebookresearch/detectron2/blob/main/detectron2/modeling/backbone/fpn.py
+        prev_features = None
+        # forward in top-down order (from low to high resolution)
+        n = len(self.convs) - 1
+        for i in range(n, -1, -1):
+            x = xs[i]
+            lateral_features = self.convs[n - i](x)
+            if i in self.fpn_top_down_levels and prev_features is not None:
+                top_down_features = F.interpolate(
+                    prev_features.to(dtype=torch.float32),
+                    scale_factor=2.0,
+                    mode=self.fpn_interp_model,
+                    align_corners=(
+                        None if self.fpn_interp_model == "nearest" else False
+                    ),
+                    antialias=False,
+                )
+                prev_features = lateral_features + top_down_features
+                if self.fuse_type == "avg":
+                    prev_features /= 2
+            else:
+                prev_features = lateral_features
+            x_out = prev_features
+            out[i] = x_out
+            pos[i] = self.position_encoding(x_out).to(x_out.dtype)
+
+        return out, pos
+
+class ViTDetNeck(nn.Module):
+    def __init__(
+        self,
+        position_encoding: nn.Module,
+        d_model: int,
+        backbone_channel_list: List[int],
+        kernel_size: int = 1,
+        stride: int = 1,
+        padding: int = 0,
+        neck_norm=None,
+    ):
+        """Initialize the neck
+        :param trunk: the backbone
+        :param position_encoding: the positional encoding to use
+        :param d_model: the dimension of the model
+        :param neck_norm: the normalization to use
+        """
+        super().__init__()
+        self.backbone_channel_list = backbone_channel_list
+        self.position_encoding = position_encoding
+        self.convs = nn.ModuleList()
+        self.d_model = d_model
+        use_bias = neck_norm is None
+        for dim in self.backbone_channel_list:
+            current = nn.Sequential()
+            current.add_module(
+                "conv_1x1",
+                nn.Conv2d(
+                    in_channels=dim,
+                    out_channels=d_model,
+                    kernel_size=1,
+                    bias=use_bias,
+                ),
+            )
+            if neck_norm is not None:
+                current.add_module("norm_0", LayerNorm2d(d_model))
+            current.add_module(
+                "conv_3x3",
+                nn.Conv2d(
+                    in_channels=d_model,
+                    out_channels=d_model,
+                    kernel_size=3,
+                    padding=1,
+                    bias=use_bias,
+                ),
+            )
+            if neck_norm is not None:
+                current.add_module("norm_1", LayerNorm2d(d_model))
+            self.convs.append(current)
+
+    def forward(self, xs: List[torch.Tensor]):
+        out = [None] * len(self.convs)
+        pos = [None] * len(self.convs)
+        assert len(xs) == len(self.convs)
+
+        x = xs[0]
+        x_out = self.convs[0](x)
+        out[0] = x_out
+        pos[0] = self.position_encoding(x_out).to(x_out.dtype)
+
+        return out, pos

utility/sam/modeling/backbones/utils.py

@@ -0,0 +1,128 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Some utilities for backbones, in particular for windowing"""
+
+import math
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def get_abs_pos(abs_pos, has_cls_token, hw):
+    """
+    Calculate absolute positional embeddings. If needed, resize embeddings and remove cls_token
+        dimension for the original embeddings.
+    Args:
+        abs_pos (Tensor): absolute positional embeddings with (1, num_position, C).
+        has_cls_token (bool): If true, has 1 embedding in abs_pos for cls token.
+        hw (Tuple): size of input image tokens.
+    Returns:
+        Absolute positional embeddings after processing with shape (1, H, W, C)
+    """
+    h, w = hw
+    if has_cls_token:
+        abs_pos = abs_pos[:, 1:]
+    xy_num = abs_pos.shape[1]
+    size = int(math.sqrt(xy_num))
+    assert size * size == xy_num
+
+    if size != h or size != w:
+        interpolate_mode = "bicubic"
+        if not torch.cuda.is_available() and torch.mps.is_available():
+            # bicubic is not supported on torch mps
+            interpolate_mode = "bilinear"
+        new_abs_pos = F.interpolate(
+            abs_pos.reshape(1, size, size, -1).permute(0, 3, 1, 2),
+            size=(h, w),
+            mode=interpolate_mode,
+            align_corners=False,
+        )
+        return new_abs_pos.permute(0, 2, 3, 1)
+    else:
+        return abs_pos.reshape(1, h, w, -1)
+
+
+def window_partition(x, window_size):
+    """
+    Partition into non-overlapping windows with padding if needed.
+    Args:
+        x (tensor): input tokens with [B, H, W, C].
+        window_size (int): window size.
+    Returns:
+        windows: windows after partition with [B * num_windows, window_size, window_size, C].
+        (Hp, Wp): padded height and width before partition
+    """
+    B, H, W, C = x.shape
+
+    pad_h = (window_size - H % window_size) % window_size
+    pad_w = (window_size - W % window_size) % window_size
+    if pad_h > 0 or pad_w > 0:
+        x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h))
+    Hp, Wp = H + pad_h, W + pad_w
+
+    x = x.view(B, Hp // window_size, window_size, Wp // window_size, window_size, C)
+    windows = x.permute(0, 1, 3, 2, 4, 5).reshape(-1, window_size, window_size, C)
+    return windows, (Hp, Wp)
+
+
+def window_unpartition(windows, window_size, pad_hw, hw):
+    """
+    Window unpartition into original sequences and removing padding.
+    Args:
+        x (tensor): input tokens with [B * num_windows, window_size, window_size, C].
+        window_size (int): window size.
+        pad_hw (Tuple): padded height and width (Hp, Wp).
+        hw (Tuple): original height and width (H, W) before padding.
+    Returns:
+        x: unpartitioned sequences with [B, H, W, C].
+    """
+    Hp, Wp = pad_hw
+    H, W = hw
+    B = windows.shape[0] // (Hp * Wp // window_size // window_size)
+    x = windows.reshape(
+        B, Hp // window_size, Wp // window_size, window_size, window_size, -1
+    )
+    x = x.permute(0, 1, 3, 2, 4, 5).reshape(B, Hp, Wp, -1)
+
+    if Hp > H or Wp > W:
+        x = x[:, :H, :W, :]
+    return x
+
+
+class PatchEmbed(nn.Module):
+    """
+    Image to Patch Embedding.
+    """
+
+    def __init__(
+        self,
+        kernel_size: Tuple[int, ...] = (7, 7),
+        stride: Tuple[int, ...] = (4, 4),
+        padding: Tuple[int, ...] = (3, 3),
+        in_chans: int = 3,
+        embed_dim: int = 768,
+    ):
+        """
+        Args:
+            kernel_size (Tuple): kernel size of the projection layer.
+            stride (Tuple): stride of the projection layer.
+            padding (Tuple): padding size of the projection layer.
+            in_chans (int): Number of input image channels.
+            embed_dim (int):  embed_dim (int): Patch embedding dimension.
+        """
+        super().__init__()
+        self.proj = nn.Conv2d(
+            in_chans, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.proj(x)
+        # B C H W -> B H W C
+        x = x.permute(0, 2, 3, 1)
+        return x

utility/sam/modeling/backbones/vitdet.py

@@ -0,0 +1,295 @@
+"""ViTDet backbone adapted from Detectron2"""
+
+from functools import partial
+from typing import List, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from sam2.modeling.backbones.utils import (
+    get_abs_pos,
+    PatchEmbed,
+    window_partition,
+    window_unpartition,
+)
+
+from sam2.modeling.sam2_utils import DropPath, LayerScale, MLP
+
+
+class Attention(nn.Module):
+    """Multi-head Attention block with relative position embeddings."""
+
+    def __init__(
+        self,
+        dim,
+        num_heads=8,
+        qkv_bias=True,
+        use_rel_pos=False,
+        rel_pos_zero_init=True,
+        input_size=None,
+    ):
+        """
+        Args:
+            dim (int): Number of input channels.
+            num_heads (int): Number of attention heads.
+            qkv_bias (bool:  If True, add a learnable bias to query, key, value.
+            rel_pos (bool): If True, add relative positional embeddings to the attention map.
+            rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
+            input_size (int or None): Input resolution for calculating the relative positional
+                parameter size.
+            attn_type: Type of attention operation, e.g. "vanilla", "vanilla-xformer".
+        """
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.proj = nn.Linear(dim, dim)
+
+        self.use_rel_pos = use_rel_pos
+
+    def forward(self, x):
+        B, H, W, _ = x.shape
+        # qkv with shape (3, B, nHead, H * W, C)
+        qkv = (
+            self.qkv(x).reshape(B, H * W, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
+        )
+        # q, k, v with shape (B * nHead, H * W, C)
+        q, k, v = qkv.reshape(3, B * self.num_heads, H * W, -1).unbind(0)
+
+        q = q.view(B, self.num_heads, H * W, -1)
+        k = k.view(B, self.num_heads, H * W, -1)
+        v = v.view(B, self.num_heads, H * W, -1)
+
+        x = F.scaled_dot_product_attention(q, k, v)
+
+        x = (
+            x.view(B, self.num_heads, H, W, -1)
+            .permute(0, 2, 3, 1, 4)
+            .reshape(B, H, W, -1)
+        )
+        x = self.proj(x)
+
+        return x
+
+
+class Block(nn.Module):
+    """Transformer blocks with support of window attention"""
+
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        drop_path=0.0,
+        norm_layer=nn.LayerNorm,
+        act_layer=nn.GELU,
+        use_rel_pos=False,
+        rel_pos_zero_init=True,
+        window_size=0,
+        input_size=None,
+        dropout=0.0,
+        init_values=None,
+    ):
+        """
+        Args:
+            dim (int): Number of input channels.
+            num_heads (int): Number of attention heads in each ViT block.
+            mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+            qkv_bias (bool): If True, add a learnable bias to query, key, value.
+            drop_path (float): Stochastic depth rate.
+            norm_layer (nn.Module): Normalization layer.
+            act_layer (nn.Module): Activation layer.
+            use_rel_pos (bool): If True, add relative positional embeddings to the attention map.
+            rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
+            window_size (int): Window size for window attention blocks. If it equals 0, then not
+                use window attention.
+            input_size (int or None): Input resolution for calculating the relative positional
+                parameter size.
+            dropout (float): Dropout rate.
+        """
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            use_rel_pos=use_rel_pos,
+            rel_pos_zero_init=rel_pos_zero_init,
+            input_size=input_size if window_size == 0 else (window_size, window_size),
+        )
+        self.ls1 = (
+            LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        )
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        self.norm2 = norm_layer(dim)
+        self.mlp = MLP(
+            dim,
+            int(dim * mlp_ratio),
+            dim,
+            num_layers=2,
+            activation=act_layer,
+        )
+        self.ls2 = (
+            LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        )
+        self.dropout = nn.Dropout(dropout)
+        self.window_size = window_size
+
+    def forward(self, x):
+        shortcut = x
+        x = self.norm1(x)
+        # Window partition
+        if self.window_size > 0:
+            H, W = x.shape[1], x.shape[2]
+            x, pad_hw = window_partition(x, self.window_size)
+
+        x = self.ls1(self.attn(x))
+        # Reverse window partition
+        if self.window_size > 0:
+            x = window_unpartition(x, self.window_size, pad_hw, (H, W))
+
+        x = shortcut + self.dropout(self.drop_path(x))
+        x = x + self.dropout(self.drop_path(self.ls2(self.mlp(self.norm2(x)))))
+
+        return x
+
+
+class ViT(nn.Module):
+    """
+    This module implements Vision Transformer (ViT) backbone in :paper:`vitdet`.
+    "Exploring Plain Vision Transformer Backbones for Object Detection",
+    https://arxiv.org/abs/2203.16527
+    """
+
+    def __init__(
+        self,
+        img_size=1024,
+        patch_size=16,
+        in_chans=3,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        drop_path_rate=0.0,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        act_layer=nn.GELU,
+        use_abs_pos=True,
+        use_rel_pos=False,
+        rel_pos_zero_init=True,
+        window_size=14,
+        window_block_indexes=(0, 1, 3, 4, 6, 7, 9, 10),
+        use_act_checkpoint=False,
+        pretrain_img_size=224,
+        pretrain_use_cls_token=True,
+        dropout=0.0,
+        weights_path=None,
+        return_interm_layers=False,
+        init_values=None,
+    ):
+        """
+        Args:
+            img_size (int): Input image size. Only relevant for rel pos.
+            patch_size (int): Patch size.
+            in_chans (int): Number of input image channels.
+            embed_dim (int): Patch embedding dimension.
+            depth (int): Depth of ViT.
+            num_heads (int): Number of attention heads in each ViT block.
+            mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+            qkv_bias (bool): If True, add a learnable bias to query, key, value.
+            drop_path_rate (float): Stochastic depth rate.
+            norm_layer (nn.Module): Normalization layer.
+            act_layer (nn.Module): Activation layer.
+            use_abs_pos (bool): If True, use absolute positional embeddings.
+            use_rel_pos (bool): If True, add relative positional embeddings to the attention map.
+            rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
+            window_size (int): Window size for window attention blocks.
+            window_block_indexes (list): Indexes for blocks using window attention.
+            residual_block_indexes (list): Indexes for blocks using conv propagation.
+            use_act_checkpoint (bool): If True, use activation checkpointing.
+            pretrain_img_size (int): input image size for pretraining models.
+            pretrain_use_cls_token (bool): If True, pretrainig models use class token.
+            dropout (float): Dropout rate. Applied in residual blocks of attn, mlp and inside the mlp.
+            path (str or None): Path to the pretrained weights.
+            return_interm_layers (bool): Whether to return intermediate layers (all global attention blocks).
+            freezing (BackboneFreezingType): Type of freezing.
+        """
+        super().__init__()
+        self.pretrain_use_cls_token = pretrain_use_cls_token
+
+        self.patch_embed = PatchEmbed(
+            kernel_size=(patch_size, patch_size),
+            stride=(patch_size, patch_size),
+            padding=(0, 0),
+            in_chans=in_chans,
+            embed_dim=embed_dim,
+        )
+
+        if use_abs_pos:
+            # Initialize absolute positional embedding with pretrain image size.
+            num_patches = (pretrain_img_size // patch_size) * (
+                pretrain_img_size // patch_size
+            )
+            num_positions = (num_patches + 1) if pretrain_use_cls_token else num_patches
+            self.pos_embed = nn.Parameter(torch.zeros(1, num_positions, embed_dim))
+        else:
+            self.pos_embed = None
+
+        # stochastic depth decay rule
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]
+
+        self.blocks = nn.ModuleList()
+        self.full_attn_ids = []
+        cur_stage = 1
+        for i in range(depth):
+            block = Block(
+                dim=embed_dim,
+                num_heads=num_heads,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                drop_path=dpr[i],
+                norm_layer=norm_layer,
+                act_layer=act_layer,
+                use_rel_pos=use_rel_pos,
+                rel_pos_zero_init=rel_pos_zero_init,
+                window_size=window_size if i in window_block_indexes else 0,
+                input_size=(img_size // patch_size, img_size // patch_size),
+                dropout=dropout,
+                init_values=init_values,
+            )
+            if i not in window_block_indexes:
+                self.full_attn_ids.append(i)
+                cur_stage += 1
+
+            self.blocks.append(block)
+
+        self.return_interm_layers = return_interm_layers
+        self.channel_list = (
+            [embed_dim] * len(self.full_attn_ids)
+            if return_interm_layers
+            else [embed_dim]
+        )
+
+    def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
+
+        x = self.patch_embed(x)
+        if self.pos_embed is not None:
+            x = x + get_abs_pos(
+                self.pos_embed, self.pretrain_use_cls_token, (x.shape[1], x.shape[2])
+            )
+
+        outputs = []
+        for i, blk in enumerate(self.blocks):
+            x = blk(x)
+            if (i == self.full_attn_ids[-1]) or (
+                self.return_interm_layers and i in self.full_attn_ids
+            ):
+                feats = x.permute(0, 3, 1, 2)
+                outputs.append(feats)
+
+        return outputs

utility/sam/modeling/memory_attention.py

@@ -0,0 +1,181 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Optional
+
+import torch
+from torch import nn, Tensor
+
+from sam2.modeling.sam.transformer import (
+    EfficientRoPEAttention1,
+    EfficientRoPEAttention2,
+    RoPEAttention,
+)
+
+from sam2.modeling.sam2_utils import get_activation_fn, get_clones
+
+
+class MemoryAttentionLayer(nn.Module):
+
+    def __init__(
+        self,
+        activation: str,
+        cross_attention: nn.Module,
+        d_model: int,
+        dim_feedforward: int,
+        dropout: float,
+        pos_enc_at_attn: bool,
+        pos_enc_at_cross_attn_keys: bool,
+        pos_enc_at_cross_attn_queries: bool,
+        self_attention: nn.Module,
+    ):
+        super().__init__()
+        self.d_model = d_model
+        self.dim_feedforward = dim_feedforward
+        self.dropout_value = dropout
+        self.self_attn = self_attention
+        self.cross_attn_image = cross_attention
+
+        # Implementation of Feedforward model
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.dropout3 = nn.Dropout(dropout)
+
+        self.activation_str = activation
+        self.activation = get_activation_fn(activation)
+
+        # Where to add pos enc
+        self.pos_enc_at_attn = pos_enc_at_attn
+        self.pos_enc_at_cross_attn_queries = pos_enc_at_cross_attn_queries
+        self.pos_enc_at_cross_attn_keys = pos_enc_at_cross_attn_keys
+
+    def _forward_sa(self, tgt, query_pos):
+        # Self-Attention
+        tgt2 = self.norm1(tgt)
+        q = k = tgt2 + query_pos if self.pos_enc_at_attn else tgt2
+        tgt2 = self.self_attn(q, k, v=tgt2)
+        tgt = tgt + self.dropout1(tgt2)
+        return tgt
+
+    def _forward_ca(self, tgt, memory, query_pos, pos, num_k_exclude_rope=0):
+        kwds = {}
+        if num_k_exclude_rope > 0:
+            assert (
+                isinstance(self.cross_attn_image, RoPEAttention)
+                or isinstance(self.cross_attn_image, EfficientRoPEAttention1)
+                or isinstance(self.cross_attn_image, EfficientRoPEAttention2)
+            )
+            kwds = {"num_k_exclude_rope": num_k_exclude_rope}
+
+        # Cross-Attention
+        tgt2 = self.norm2(tgt)
+        tgt2 = self.cross_attn_image(
+            q=tgt2 + query_pos if self.pos_enc_at_cross_attn_queries else tgt2,
+            k=memory + pos if self.pos_enc_at_cross_attn_keys else memory,
+            v=memory,
+            **kwds,
+        )
+        tgt = tgt + self.dropout2(tgt2)
+        return tgt
+
+    def forward(
+        self,
+        tgt,
+        memory,
+        pos: Optional[Tensor] = None,
+        query_pos: Optional[Tensor] = None,
+        num_k_exclude_rope: int = 0,
+    ) -> torch.Tensor:
+
+        # Self-Attn, Cross-Attn
+        tgt = self._forward_sa(tgt, query_pos)
+        tgt = self._forward_ca(tgt, memory, query_pos, pos, num_k_exclude_rope)
+        # MLP
+        tgt2 = self.norm3(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
+        tgt = tgt + self.dropout3(tgt2)
+        return tgt
+
+
+class MemoryAttention(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        pos_enc_at_input: bool,
+        layer: nn.Module,
+        num_layers: int,
+        batch_first: bool = True,  # Do layers expect batch first input?
+    ):
+        super().__init__()
+        self.d_model = d_model
+        self.layers = get_clones(layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = nn.LayerNorm(d_model)
+        self.pos_enc_at_input = pos_enc_at_input
+        self.batch_first = batch_first
+
+    def forward(
+        self,
+        curr: torch.Tensor,  # self-attention inputs
+        memory: torch.Tensor,  # cross-attention inputs
+        curr_pos: Optional[Tensor] = None,  # pos_enc for self-attention inputs
+        memory_pos: Optional[Tensor] = None,  # pos_enc for cross-attention inputs
+        num_obj_ptr_tokens: int = 0,  # number of object pointer *tokens*
+    ):
+        if isinstance(curr, list):
+            assert isinstance(curr_pos, list)
+            assert len(curr) == len(curr_pos) == 1
+            curr, curr_pos = (
+                curr[0],
+                curr_pos[0],
+            )
+
+        assert (
+            curr.shape[1] == memory.shape[1]
+        ), "Batch size must be the same for curr and memory"
+
+        output = curr
+        if self.pos_enc_at_input and curr_pos is not None:
+            output = output + 0.1 * curr_pos
+
+        if self.batch_first:
+            # Convert to batch first
+            output = output.transpose(0, 1)
+            curr_pos = curr_pos.transpose(0, 1)
+            memory = memory.transpose(0, 1)
+            memory_pos = memory_pos.transpose(0, 1)
+
+        for layer in self.layers:
+            kwds = {}
+            if (
+                isinstance(layer.cross_attn_image, RoPEAttention)
+                or isinstance(layer.cross_attn_image, EfficientRoPEAttention1)
+                or isinstance(layer.cross_attn_image, EfficientRoPEAttention2)
+            ):
+                kwds = {"num_k_exclude_rope": num_obj_ptr_tokens}
+
+            output = layer(
+                tgt=output,
+                memory=memory,
+                pos=memory_pos,
+                query_pos=curr_pos,
+                **kwds,
+            )
+        normed_output = self.norm(output)
+
+        if self.batch_first:
+            # Convert back to seq first
+            normed_output = normed_output.transpose(0, 1)
+            curr_pos = curr_pos.transpose(0, 1)
+
+        return normed_output

utility/sam/modeling/memory_encoder.py

@@ -0,0 +1,181 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from sam2.modeling.sam2_utils import DropPath, get_clones, LayerNorm2d
+
+
+class MaskDownSampler(nn.Module):
+    """
+    Progressively downsample a mask by total_stride, each time by stride.
+    Note that LayerNorm is applied per *token*, like in ViT.
+
+    With each downsample (by a factor stride**2), channel capacity increases by the same factor.
+    In the end, we linearly project to embed_dim channels.
+    """
+
+    def __init__(
+        self,
+        embed_dim=256,
+        kernel_size=4,
+        stride=4,
+        padding=0,
+        total_stride=16,
+        activation=nn.GELU,
+    ):
+        super().__init__()
+        num_layers = int(math.log2(total_stride) // math.log2(stride))
+        assert stride**num_layers == total_stride
+        self.encoder = nn.Sequential()
+        mask_in_chans, mask_out_chans = 1, 1
+        for _ in range(num_layers):
+            mask_out_chans = mask_in_chans * (stride**2)
+            self.encoder.append(
+                nn.Conv2d(
+                    mask_in_chans,
+                    mask_out_chans,
+                    kernel_size=kernel_size,
+                    stride=stride,
+                    padding=padding,
+                )
+            )
+            self.encoder.append(LayerNorm2d(mask_out_chans))
+            self.encoder.append(activation())
+            mask_in_chans = mask_out_chans
+
+        self.encoder.append(nn.Conv2d(mask_out_chans, embed_dim, kernel_size=1))
+
+    def forward(self, x):
+        return self.encoder(x)
+
+
+# Lightly adapted from ConvNext (https://github.com/facebookresearch/ConvNeXt)
+class CXBlock(nn.Module):
+    r"""ConvNeXt Block. There are two equivalent implementations:
+    (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)
+    (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back
+    We use (2) as we find it slightly faster in PyTorch
+
+    Args:
+        dim (int): Number of input channels.
+        drop_path (float): Stochastic depth rate. Default: 0.0
+        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
+    """
+
+    def __init__(
+        self,
+        dim,
+        kernel_size=7,
+        padding=3,
+        drop_path=0.0,
+        layer_scale_init_value=1e-6,
+        use_dwconv=True,
+    ):
+        super().__init__()
+        self.dwconv = nn.Conv2d(
+            dim,
+            dim,
+            kernel_size=kernel_size,
+            padding=padding,
+            groups=dim if use_dwconv else 1,
+        )  # depthwise conv
+        self.norm = LayerNorm2d(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(
+            dim, 4 * dim
+        )  # pointwise/1x1 convs, implemented with linear layers
+        self.act = nn.GELU()
+        self.pwconv2 = nn.Linear(4 * dim, dim)
+        self.gamma = (
+            nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True)
+            if layer_scale_init_value > 0
+            else None
+        )
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+    def forward(self, x):
+        input = x
+        x = self.dwconv(x)
+        x = self.norm(x)
+        x = x.permute(0, 2, 3, 1)  # (N, C, H, W) -> (N, H, W, C)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.pwconv2(x)
+        if self.gamma is not None:
+            x = self.gamma * x
+        x = x.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)
+
+        x = input + self.drop_path(x)
+        return x
+
+
+class Fuser(nn.Module):
+    def __init__(self, layer, num_layers, dim=None, input_projection=False):
+        super().__init__()
+        self.proj = nn.Identity()
+        self.layers = get_clones(layer, num_layers)
+
+        if input_projection:
+            assert dim is not None
+            self.proj = nn.Conv2d(dim, dim, kernel_size=1)
+
+    def forward(self, x):
+        # normally x: (N, C, H, W)
+        x = self.proj(x)
+        for layer in self.layers:
+            x = layer(x)
+        return x
+
+
+class MemoryEncoder(nn.Module):
+    def __init__(
+        self,
+        out_dim,
+        mask_downsampler,
+        fuser,
+        position_encoding,
+        in_dim=256,  # in_dim of pix_feats
+    ):
+        super().__init__()
+
+        self.mask_downsampler = mask_downsampler
+
+        self.pix_feat_proj = nn.Conv2d(in_dim, in_dim, kernel_size=1)
+        self.fuser = fuser
+        self.position_encoding = position_encoding
+        self.out_proj = nn.Identity()
+        if out_dim != in_dim:
+            self.out_proj = nn.Conv2d(in_dim, out_dim, kernel_size=1)
+
+    def forward(
+        self,
+        pix_feat: torch.Tensor,
+        masks: torch.Tensor,
+        skip_mask_sigmoid: bool = False,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        ## Process masks
+        # sigmoid, so that less domain shift from gt masks which are bool
+        if not skip_mask_sigmoid:
+            masks = F.sigmoid(masks)
+        masks = self.mask_downsampler(masks)
+
+        ## Fuse pix_feats and downsampled masks
+        # in case the visual features are on CPU, cast them to CUDA
+        pix_feat = pix_feat.to(masks.device)
+
+        x = self.pix_feat_proj(pix_feat)
+        x = x + masks
+        x = self.fuser(x)
+        x = self.out_proj(x)
+
+        pos = self.position_encoding(x).to(x.dtype)
+
+        return {"vision_features": x, "vision_pos_enc": [pos]}

utility/sam/modeling/position_encoding.py

@@ -0,0 +1,239 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+from typing import Any, Optional, Tuple
+
+import numpy as np
+
+import torch
+from torch import nn
+
+
+class PositionEmbeddingSine(nn.Module):
+    """
+    This is a more standard version of the position embedding, very similar to the one
+    used by the Attention Is All You Need paper, generalized to work on images.
+    """
+
+    def __init__(
+        self,
+        num_pos_feats,
+        temperature: int = 10000,
+        normalize: bool = True,
+        scale: Optional[float] = None,
+        # Following settings only relevant
+        # for warmping up cache for compilation
+        warmup_cache: bool = True,
+        image_size: int = 1024,
+        strides: Tuple[int] = (4, 8, 16, 32),
+    ):
+        super().__init__()
+        assert num_pos_feats % 2 == 0, "Expecting even model width"
+        self.num_pos_feats = num_pos_feats // 2
+        self.temperature = temperature
+        self.normalize = normalize
+        if scale is not None and normalize is False:
+            raise ValueError("normalize should be True if scale is passed")
+        if scale is None:
+            scale = 2 * math.pi
+        self.scale = scale
+
+        self.cache = {}
+        if warmup_cache and torch.cuda.is_available():
+            # Warmup cache for cuda, to help with compilation
+            device = torch.device("cuda")
+            for stride in strides:
+                cache_key = (image_size // stride, image_size // stride)
+                self._pe(1, device, *cache_key)
+
+    def _encode_xy(self, x, y):
+        # The positions are expected to be normalized
+        assert len(x) == len(y) and x.ndim == y.ndim == 1
+        x_embed = x * self.scale
+        y_embed = y * self.scale
+
+        dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)
+        dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)
+
+        pos_x = x_embed[:, None] / dim_t
+        pos_y = y_embed[:, None] / dim_t
+        pos_x = torch.stack(
+            (pos_x[:, 0::2].sin(), pos_x[:, 1::2].cos()), dim=2
+        ).flatten(1)
+        pos_y = torch.stack(
+            (pos_y[:, 0::2].sin(), pos_y[:, 1::2].cos()), dim=2
+        ).flatten(1)
+        return pos_x, pos_y
+
+    @torch.no_grad()
+    def encode_boxes(self, x, y, w, h):
+        pos_x, pos_y = self._encode_xy(x, y)
+        pos = torch.cat((pos_y, pos_x, h[:, None], w[:, None]), dim=1)
+        return pos
+
+    encode = encode_boxes  # Backwards compatibility
+
+    @torch.no_grad()
+    def encode_points(self, x, y, labels):
+        (bx, nx), (by, ny), (bl, nl) = x.shape, y.shape, labels.shape
+        assert bx == by and nx == ny and bx == bl and nx == nl
+        pos_x, pos_y = self._encode_xy(x.flatten(), y.flatten())
+        pos_x, pos_y = pos_x.reshape(bx, nx, -1), pos_y.reshape(by, ny, -1)
+        pos = torch.cat((pos_y, pos_x, labels[:, :, None]), dim=2)
+        return pos
+
+    @torch.no_grad()
+    def _pe(self, B, device, *cache_key):
+        H, W = cache_key
+        if cache_key in self.cache:
+            return self.cache[cache_key].to(device)[None].repeat(B, 1, 1, 1)
+
+        y_embed = (
+            torch.arange(1, H + 1, dtype=torch.float32, device=device)
+            .view(1, -1, 1)
+            .repeat(B, 1, W)
+        )
+        x_embed = (
+            torch.arange(1, W + 1, dtype=torch.float32, device=device)
+            .view(1, 1, -1)
+            .repeat(B, H, 1)
+        )
+
+        if self.normalize:
+            eps = 1e-6
+            y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
+            x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale
+
+        dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=device)
+        dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)
+
+        pos_x = x_embed[:, :, :, None] / dim_t
+        pos_y = y_embed[:, :, :, None] / dim_t
+        pos_x = torch.stack(
+            (pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4
+        ).flatten(3)
+        pos_y = torch.stack(
+            (pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4
+        ).flatten(3)
+        pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
+        self.cache[cache_key] = pos[0]
+        return pos
+
+    @torch.no_grad()
+    def forward(self, x: torch.Tensor):
+        B = x.shape[0]
+        cache_key = (x.shape[-2], x.shape[-1])
+        return self._pe(B, x.device, *cache_key)
+
+
+class PositionEmbeddingRandom(nn.Module):
+    """
+    Positional encoding using random spatial frequencies.
+    """
+
+    def __init__(self, num_pos_feats: int = 64, scale: Optional[float] = None) -> None:
+        super().__init__()
+        if scale is None or scale <= 0.0:
+            scale = 1.0
+        self.register_buffer(
+            "positional_encoding_gaussian_matrix",
+            scale * torch.randn((2, num_pos_feats)),
+        )
+
+    def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor:
+        """Positionally encode points that are normalized to [0,1]."""
+        # assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape
+        coords = 2 * coords - 1
+        coords = coords @ self.positional_encoding_gaussian_matrix
+        coords = 2 * np.pi * coords
+        # outputs d_1 x ... x d_n x C shape
+        return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1)
+
+    def forward(self, size: Tuple[int, int]) -> torch.Tensor:
+        """Generate positional encoding for a grid of the specified size."""
+        h, w = size
+        device: Any = self.positional_encoding_gaussian_matrix.device
+        grid = torch.ones((h, w), device=device, dtype=torch.float32)
+        y_embed = grid.cumsum(dim=0) - 0.5
+        x_embed = grid.cumsum(dim=1) - 0.5
+        y_embed = y_embed / h
+        x_embed = x_embed / w
+
+        pe = self._pe_encoding(torch.stack([x_embed, y_embed], dim=-1))
+        return pe.permute(2, 0, 1)  # C x H x W
+
+    def forward_with_coords(
+        self, coords_input: torch.Tensor, image_size: Tuple[int, int]
+    ) -> torch.Tensor:
+        """Positionally encode points that are not normalized to [0,1]."""
+        coords = coords_input.clone()
+        coords[:, :, 0] = coords[:, :, 0] / image_size[1]
+        coords[:, :, 1] = coords[:, :, 1] / image_size[0]
+        return self._pe_encoding(coords.to(torch.float))  # B x N x C
+
+
+# Rotary Positional Encoding, adapted from:
+# 1. https://github.com/meta-llama/codellama/blob/main/llama/model.py
+# 2. https://github.com/naver-ai/rope-vit
+# 3. https://github.com/lucidrains/rotary-embedding-torch
+
+
+def init_t_xy(end_x: int, end_y: int):
+    t = torch.arange(end_x * end_y, dtype=torch.float32)
+    t_x = (t % end_x).float()
+    t_y = torch.div(t, end_x, rounding_mode="floor").float()
+    return t_x, t_y
+
+
+def compute_axial_cis(dim: int, end_x: int, end_y: int, theta: float = 10000.0):
+    freqs_x = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim))
+    freqs_y = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim))
+
+    t_x, t_y = init_t_xy(end_x, end_y)
+    freqs_x = torch.outer(t_x, freqs_x)
+    freqs_y = torch.outer(t_y, freqs_y)
+    freqs_cis_x = torch.polar(torch.ones_like(freqs_x), freqs_x)
+    freqs_cis_y = torch.polar(torch.ones_like(freqs_y), freqs_y)
+    return torch.cat([freqs_cis_x, freqs_cis_y], dim=-1)
+
+
+def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
+    ndim = x.ndim
+    assert 0 <= 1 < ndim
+    assert freqs_cis.shape == (x.shape[-2], x.shape[-1])
+    shape = [d if i >= ndim - 2 else 1 for i, d in enumerate(x.shape)]
+    return freqs_cis.view(*shape)
+
+
+def apply_rotary_enc(
+    xq: torch.Tensor,
+    xk: torch.Tensor,
+    freqs_cis: torch.Tensor,
+    repeat_freqs_k: bool = False,
+):
+    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
+    xk_ = (
+        torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
+        if xk.shape[-2] != 0
+        else None
+    )
+    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
+    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
+    if xk_ is None:
+        # no keys to rotate, due to dropout
+        return xq_out.type_as(xq).to(xq.device), xk
+    # repeat freqs along seq_len dim to match k seq_len
+    if repeat_freqs_k:
+        r = xk_.shape[-2] // xq_.shape[-2]
+        if freqs_cis.is_cuda:
+            freqs_cis = freqs_cis.repeat(*([1] * (freqs_cis.ndim - 2)), r, 1)
+        else:
+            # torch.repeat on complex numbers may not be supported on non-CUDA devices
+            # (freqs_cis has 4 dims and we repeat on dim 2) so we use expand + flatten
+            freqs_cis = freqs_cis.unsqueeze(2).expand(-1, -1, r, -1, -1).flatten(2, 3)
+    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
+    return xq_out.type_as(xq).to(xq.device), xk_out.type_as(xk).to(xk.device)

utility/sam/modeling/sam/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.

utility/sam/modeling/sam/mask_decoder.py

@@ -0,0 +1,295 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import List, Optional, Tuple, Type
+
+import torch
+from torch import nn
+
+from sam2.modeling.sam2_utils import LayerNorm2d, MLP
+
+
+class MaskDecoder(nn.Module):
+    def __init__(
+        self,
+        *,
+        transformer_dim: int,
+        transformer: nn.Module,
+        num_multimask_outputs: int = 3,
+        activation: Type[nn.Module] = nn.GELU,
+        iou_head_depth: int = 3,
+        iou_head_hidden_dim: int = 256,
+        use_high_res_features: bool = False,
+        iou_prediction_use_sigmoid=False,
+        dynamic_multimask_via_stability=False,
+        dynamic_multimask_stability_delta=0.05,
+        dynamic_multimask_stability_thresh=0.98,
+        pred_obj_scores: bool = False,
+        pred_obj_scores_mlp: bool = False,
+        use_multimask_token_for_obj_ptr: bool = False,
+    ) -> None:
+        """
+        Predicts masks given an image and prompt embeddings, using a
+        transformer architecture.
+
+        Arguments:
+          transformer_dim (int): the channel dimension of the transformer
+          transformer (nn.Module): the transformer used to predict masks
+          num_multimask_outputs (int): the number of masks to predict
+            when disambiguating masks
+          activation (nn.Module): the type of activation to use when
+            upscaling masks
+          iou_head_depth (int): the depth of the MLP used to predict
+            mask quality
+          iou_head_hidden_dim (int): the hidden dimension of the MLP
+            used to predict mask quality
+        """
+        super().__init__()
+        self.transformer_dim = transformer_dim
+        self.transformer = transformer
+
+        self.num_multimask_outputs = num_multimask_outputs
+
+        self.iou_token = nn.Embedding(1, transformer_dim)
+        self.num_mask_tokens = num_multimask_outputs + 1
+        self.mask_tokens = nn.Embedding(self.num_mask_tokens, transformer_dim)
+
+        self.pred_obj_scores = pred_obj_scores
+        if self.pred_obj_scores:
+            self.obj_score_token = nn.Embedding(1, transformer_dim)
+        self.use_multimask_token_for_obj_ptr = use_multimask_token_for_obj_ptr
+
+        self.output_upscaling = nn.Sequential(
+            nn.ConvTranspose2d(
+                transformer_dim, transformer_dim // 4, kernel_size=2, stride=2
+            ),
+            LayerNorm2d(transformer_dim // 4),
+            activation(),
+            nn.ConvTranspose2d(
+                transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2
+            ),
+            activation(),
+        )
+        self.use_high_res_features = use_high_res_features
+        if use_high_res_features:
+            self.conv_s0 = nn.Conv2d(
+                transformer_dim, transformer_dim // 8, kernel_size=1, stride=1
+            )
+            self.conv_s1 = nn.Conv2d(
+                transformer_dim, transformer_dim // 4, kernel_size=1, stride=1
+            )
+
+        self.output_hypernetworks_mlps = nn.ModuleList(
+            [
+                MLP(transformer_dim, transformer_dim, transformer_dim // 8, 3)
+                for i in range(self.num_mask_tokens)
+            ]
+        )
+
+        self.iou_prediction_head = MLP(
+            transformer_dim,
+            iou_head_hidden_dim,
+            self.num_mask_tokens,
+            iou_head_depth,
+            sigmoid_output=iou_prediction_use_sigmoid,
+        )
+        if self.pred_obj_scores:
+            self.pred_obj_score_head = nn.Linear(transformer_dim, 1)
+            if pred_obj_scores_mlp:
+                self.pred_obj_score_head = MLP(transformer_dim, transformer_dim, 1, 3)
+
+        # When outputting a single mask, optionally we can dynamically fall back to the best
+        # multimask output token if the single mask output token gives low stability scores.
+        self.dynamic_multimask_via_stability = dynamic_multimask_via_stability
+        self.dynamic_multimask_stability_delta = dynamic_multimask_stability_delta
+        self.dynamic_multimask_stability_thresh = dynamic_multimask_stability_thresh
+
+    def forward(
+        self,
+        image_embeddings: torch.Tensor,
+        image_pe: torch.Tensor,
+        sparse_prompt_embeddings: torch.Tensor,
+        dense_prompt_embeddings: torch.Tensor,
+        multimask_output: bool,
+        repeat_image: bool,
+        high_res_features: Optional[List[torch.Tensor]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Predict masks given image and prompt embeddings.
+
+        Arguments:
+          image_embeddings (torch.Tensor): the embeddings from the image encoder
+          image_pe (torch.Tensor): positional encoding with the shape of image_embeddings
+          sparse_prompt_embeddings (torch.Tensor): the embeddings of the points and boxes
+          dense_prompt_embeddings (torch.Tensor): the embeddings of the mask inputs
+          multimask_output (bool): Whether to return multiple masks or a single
+            mask.
+
+        Returns:
+          torch.Tensor: batched predicted masks
+          torch.Tensor: batched predictions of mask quality
+          torch.Tensor: batched SAM token for mask output
+        """
+        masks, iou_pred, mask_tokens_out, object_score_logits = self.predict_masks(
+            image_embeddings=image_embeddings,
+            image_pe=image_pe,
+            sparse_prompt_embeddings=sparse_prompt_embeddings,
+            dense_prompt_embeddings=dense_prompt_embeddings,
+            repeat_image=repeat_image,
+            high_res_features=high_res_features,
+        )
+
+        # Select the correct mask or masks for output
+        if multimask_output:
+            masks = masks[:, 1:, :, :]
+            iou_pred = iou_pred[:, 1:]
+        elif self.dynamic_multimask_via_stability and not self.training:
+            masks, iou_pred = self._dynamic_multimask_via_stability(masks, iou_pred)
+        else:
+            masks = masks[:, 0:1, :, :]
+            iou_pred = iou_pred[:, 0:1]
+
+        if multimask_output and self.use_multimask_token_for_obj_ptr:
+            sam_tokens_out = mask_tokens_out[:, 1:]  # [b, 3, c] shape
+        else:
+            # Take the mask output token. Here we *always* use the token for single mask output.
+            # At test time, even if we track after 1-click (and using multimask_output=True),
+            # we still take the single mask token here. The rationale is that we always track
+            # after multiple clicks during training, so the past tokens seen during training
+            # are always the single mask token (and we'll let it be the object-memory token).
+            sam_tokens_out = mask_tokens_out[:, 0:1]  # [b, 1, c] shape
+
+        # Prepare output
+        return masks, iou_pred, sam_tokens_out, object_score_logits
+
+    def predict_masks(
+        self,
+        image_embeddings: torch.Tensor,
+        image_pe: torch.Tensor,
+        sparse_prompt_embeddings: torch.Tensor,
+        dense_prompt_embeddings: torch.Tensor,
+        repeat_image: bool,
+        high_res_features: Optional[List[torch.Tensor]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Predicts masks. See 'forward' for more details."""
+        # Concatenate output tokens
+        s = 0
+        if self.pred_obj_scores:
+            output_tokens = torch.cat(
+                [
+                    self.obj_score_token.weight,
+                    self.iou_token.weight,
+                    self.mask_tokens.weight,
+                ],
+                dim=0,
+            )
+            s = 1
+        else:
+            output_tokens = torch.cat(
+                [self.iou_token.weight, self.mask_tokens.weight], dim=0
+            )
+        output_tokens = output_tokens.unsqueeze(0).expand(
+            sparse_prompt_embeddings.size(0), -1, -1
+        )
+        tokens = torch.cat((output_tokens, sparse_prompt_embeddings), dim=1)
+
+        # Expand per-image data in batch direction to be per-mask
+        if repeat_image:
+            src = torch.repeat_interleave(image_embeddings, tokens.shape[0], dim=0)
+        else:
+            assert image_embeddings.shape[0] == tokens.shape[0]
+            src = image_embeddings
+        src = src + dense_prompt_embeddings
+        assert (
+            image_pe.size(0) == 1
+        ), "image_pe should have size 1 in batch dim (from `get_dense_pe()`)"
+        pos_src = torch.repeat_interleave(image_pe, tokens.shape[0], dim=0)
+        b, c, h, w = src.shape
+
+        # Run the transformer
+        hs, src = self.transformer(src, pos_src, tokens)
+        iou_token_out = hs[:, s, :]
+        mask_tokens_out = hs[:, s + 1 : (s + 1 + self.num_mask_tokens), :]
+
+        # Upscale mask embeddings and predict masks using the mask tokens
+        src = src.transpose(1, 2).view(b, c, h, w)
+        if not self.use_high_res_features:
+            upscaled_embedding = self.output_upscaling(src)
+        else:
+            dc1, ln1, act1, dc2, act2 = self.output_upscaling
+            feat_s0, feat_s1 = high_res_features
+            upscaled_embedding = act1(ln1(dc1(src) + feat_s1))
+            upscaled_embedding = act2(dc2(upscaled_embedding) + feat_s0)
+
+        hyper_in_list: List[torch.Tensor] = []
+        for i in range(self.num_mask_tokens):
+            hyper_in_list.append(
+                self.output_hypernetworks_mlps[i](mask_tokens_out[:, i, :])
+            )
+        hyper_in = torch.stack(hyper_in_list, dim=1)
+        b, c, h, w = upscaled_embedding.shape
+        masks = (hyper_in @ upscaled_embedding.view(b, c, h * w)).view(b, -1, h, w)
+
+        # Generate mask quality predictions
+        iou_pred = self.iou_prediction_head(iou_token_out)
+        if self.pred_obj_scores:
+            assert s == 1
+            object_score_logits = self.pred_obj_score_head(hs[:, 0, :])
+        else:
+            # Obj scores logits - default to 10.0, i.e. assuming the object is present, sigmoid(10)=1
+            object_score_logits = 10.0 * iou_pred.new_ones(iou_pred.shape[0], 1)
+
+        return masks, iou_pred, mask_tokens_out, object_score_logits
+
+    def _get_stability_scores(self, mask_logits):
+        """
+        Compute stability scores of the mask logits based on the IoU between upper and
+        lower thresholds.
+        """
+        mask_logits = mask_logits.flatten(-2)
+        stability_delta = self.dynamic_multimask_stability_delta
+        area_i = torch.sum(mask_logits > stability_delta, dim=-1).float()
+        area_u = torch.sum(mask_logits > -stability_delta, dim=-1).float()
+        stability_scores = torch.where(area_u > 0, area_i / area_u, 1.0)
+        return stability_scores
+
+    def _dynamic_multimask_via_stability(self, all_mask_logits, all_iou_scores):
+        """
+        When outputting a single mask, if the stability score from the current single-mask
+        output (based on output token 0) falls below a threshold, we instead select from
+        multi-mask outputs (based on output token 1~3) the mask with the highest predicted
+        IoU score. This is intended to ensure a valid mask for both clicking and tracking.
+        """
+        # The best mask from multimask output tokens (1~3)
+        multimask_logits = all_mask_logits[:, 1:, :, :]
+        multimask_iou_scores = all_iou_scores[:, 1:]
+        best_scores_inds = torch.argmax(multimask_iou_scores, dim=-1)
+        batch_inds = torch.arange(
+            multimask_iou_scores.size(0), device=all_iou_scores.device
+        )
+        best_multimask_logits = multimask_logits[batch_inds, best_scores_inds]
+        best_multimask_logits = best_multimask_logits.unsqueeze(1)
+        best_multimask_iou_scores = multimask_iou_scores[batch_inds, best_scores_inds]
+        best_multimask_iou_scores = best_multimask_iou_scores.unsqueeze(1)
+
+        # The mask from singlemask output token 0 and its stability score
+        singlemask_logits = all_mask_logits[:, 0:1, :, :]
+        singlemask_iou_scores = all_iou_scores[:, 0:1]
+        stability_scores = self._get_stability_scores(singlemask_logits)
+        is_stable = stability_scores >= self.dynamic_multimask_stability_thresh
+
+        # Dynamically fall back to best multimask output upon low stability scores.
+        mask_logits_out = torch.where(
+            is_stable[..., None, None].expand_as(singlemask_logits),
+            singlemask_logits,
+            best_multimask_logits,
+        )
+        iou_scores_out = torch.where(
+            is_stable.expand_as(singlemask_iou_scores),
+            singlemask_iou_scores,
+            best_multimask_iou_scores,
+        )
+        return mask_logits_out, iou_scores_out

utility/sam/modeling/sam/prompt_encoder.py

@@ -0,0 +1,202 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Optional, Tuple, Type
+
+import torch
+from torch import nn
+
+from sam2.modeling.position_encoding import PositionEmbeddingRandom
+
+from sam2.modeling.sam2_utils import LayerNorm2d
+
+
+class PromptEncoder(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        image_embedding_size: Tuple[int, int],
+        input_image_size: Tuple[int, int],
+        mask_in_chans: int,
+        activation: Type[nn.Module] = nn.GELU,
+    ) -> None:
+        """
+        Encodes prompts for input to SAM's mask decoder.
+
+        Arguments:
+          embed_dim (int): The prompts' embedding dimension
+          image_embedding_size (tuple(int, int)): The spatial size of the
+            image embedding, as (H, W).
+          input_image_size (int): The padded size of the image as input
+            to the image encoder, as (H, W).
+          mask_in_chans (int): The number of hidden channels used for
+            encoding input masks.
+          activation (nn.Module): The activation to use when encoding
+            input masks.
+        """
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.input_image_size = input_image_size
+        self.image_embedding_size = image_embedding_size
+        self.pe_layer = PositionEmbeddingRandom(embed_dim // 2)
+
+        self.num_point_embeddings: int = 4  # pos/neg point + 2 box corners
+        point_embeddings = [
+            nn.Embedding(1, embed_dim) for i in range(self.num_point_embeddings)
+        ]
+        self.point_embeddings = nn.ModuleList(point_embeddings)
+        self.not_a_point_embed = nn.Embedding(1, embed_dim)
+
+        self.mask_input_size = (
+            4 * image_embedding_size[0],
+            4 * image_embedding_size[1],
+        )
+        self.mask_downscaling = nn.Sequential(
+            nn.Conv2d(1, mask_in_chans // 4, kernel_size=2, stride=2),
+            LayerNorm2d(mask_in_chans // 4),
+            activation(),
+            nn.Conv2d(mask_in_chans // 4, mask_in_chans, kernel_size=2, stride=2),
+            LayerNorm2d(mask_in_chans),
+            activation(),
+            nn.Conv2d(mask_in_chans, embed_dim, kernel_size=1),
+        )
+        self.no_mask_embed = nn.Embedding(1, embed_dim)
+
+    def get_dense_pe(self) -> torch.Tensor:
+        """
+        Returns the positional encoding used to encode point prompts,
+        applied to a dense set of points the shape of the image encoding.
+
+        Returns:
+          torch.Tensor: Positional encoding with shape
+            1x(embed_dim)x(embedding_h)x(embedding_w)
+        """
+        return self.pe_layer(self.image_embedding_size).unsqueeze(0)
+
+    def _embed_points(
+        self,
+        points: torch.Tensor,
+        labels: torch.Tensor,
+        pad: bool,
+    ) -> torch.Tensor:
+        """Embeds point prompts."""
+        points = points + 0.5  # Shift to center of pixel
+        if pad:
+            padding_point = torch.zeros((points.shape[0], 1, 2), device=points.device)
+            padding_label = -torch.ones((labels.shape[0], 1), device=labels.device)
+            points = torch.cat([points, padding_point], dim=1)
+            labels = torch.cat([labels, padding_label], dim=1)
+        point_embedding = self.pe_layer.forward_with_coords(
+            points, self.input_image_size
+        )
+
+        point_embedding = torch.where(
+            (labels == -1).unsqueeze(-1),
+            torch.zeros_like(point_embedding) + self.not_a_point_embed.weight,
+            point_embedding,
+        )
+        point_embedding = torch.where(
+            (labels == 0).unsqueeze(-1),
+            point_embedding + self.point_embeddings[0].weight,
+            point_embedding,
+        )
+        point_embedding = torch.where(
+            (labels == 1).unsqueeze(-1),
+            point_embedding + self.point_embeddings[1].weight,
+            point_embedding,
+        )
+        point_embedding = torch.where(
+            (labels == 2).unsqueeze(-1),
+            point_embedding + self.point_embeddings[2].weight,
+            point_embedding,
+        )
+        point_embedding = torch.where(
+            (labels == 3).unsqueeze(-1),
+            point_embedding + self.point_embeddings[3].weight,
+            point_embedding,
+        )
+        return point_embedding
+
+    def _embed_boxes(self, boxes: torch.Tensor) -> torch.Tensor:
+        """Embeds box prompts."""
+        boxes = boxes + 0.5  # Shift to center of pixel
+        coords = boxes.reshape(-1, 2, 2)
+        corner_embedding = self.pe_layer.forward_with_coords(
+            coords, self.input_image_size
+        )
+        corner_embedding[:, 0, :] += self.point_embeddings[2].weight
+        corner_embedding[:, 1, :] += self.point_embeddings[3].weight
+        return corner_embedding
+
+    def _embed_masks(self, masks: torch.Tensor) -> torch.Tensor:
+        """Embeds mask inputs."""
+        mask_embedding = self.mask_downscaling(masks)
+        return mask_embedding
+
+    def _get_batch_size(
+        self,
+        points: Optional[Tuple[torch.Tensor, torch.Tensor]],
+        boxes: Optional[torch.Tensor],
+        masks: Optional[torch.Tensor],
+    ) -> int:
+        """
+        Gets the batch size of the output given the batch size of the input prompts.
+        """
+        if points is not None:
+            return points[0].shape[0]
+        elif boxes is not None:
+            return boxes.shape[0]
+        elif masks is not None:
+            return masks.shape[0]
+        else:
+            return 1
+
+    def _get_device(self) -> torch.device:
+        return self.point_embeddings[0].weight.device
+
+    def forward(
+        self,
+        points: Optional[Tuple[torch.Tensor, torch.Tensor]],
+        boxes: Optional[torch.Tensor],
+        masks: Optional[torch.Tensor],
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Embeds different types of prompts, returning both sparse and dense
+        embeddings.
+
+        Arguments:
+          points (tuple(torch.Tensor, torch.Tensor) or none): point coordinates
+            and labels to embed.
+          boxes (torch.Tensor or none): boxes to embed
+          masks (torch.Tensor or none): masks to embed
+
+        Returns:
+          torch.Tensor: sparse embeddings for the points and boxes, with shape
+            BxNx(embed_dim), where N is determined by the number of input points
+            and boxes.
+          torch.Tensor: dense embeddings for the masks, in the shape
+            Bx(embed_dim)x(embed_H)x(embed_W)
+        """
+        bs = self._get_batch_size(points, boxes, masks)
+        sparse_embeddings = torch.empty(
+            (bs, 0, self.embed_dim), device=self._get_device()
+        )
+        if points is not None:
+            coords, labels = points
+            point_embeddings = self._embed_points(coords, labels, pad=(boxes is None))
+            sparse_embeddings = torch.cat([sparse_embeddings, point_embeddings], dim=1)
+        if boxes is not None:
+            box_embeddings = self._embed_boxes(boxes)
+            sparse_embeddings = torch.cat([sparse_embeddings, box_embeddings], dim=1)
+
+        if masks is not None:
+            dense_embeddings = self._embed_masks(masks)
+        else:
+            dense_embeddings = self.no_mask_embed.weight.reshape(1, -1, 1, 1).expand(
+                bs, -1, self.image_embedding_size[0], self.image_embedding_size[1]
+            )
+
+        return sparse_embeddings, dense_embeddings

utility/sam/modeling/sam/transformer.py

@@ -0,0 +1,536 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+from functools import partial
+from typing import Tuple, Type
+
+import torch
+import torch.nn.functional as F
+from torch import nn, Tensor
+
+from sam2.modeling.position_encoding import apply_rotary_enc, compute_axial_cis
+from sam2.modeling.sam2_utils import MLP
+
+
+class TwoWayTransformer(nn.Module):
+    def __init__(
+        self,
+        depth: int,
+        embedding_dim: int,
+        num_heads: int,
+        mlp_dim: int,
+        activation: Type[nn.Module] = nn.ReLU,
+        attention_downsample_rate: int = 2,
+    ) -> None:
+        """
+        A transformer decoder that attends to an input image using
+        queries whose positional embedding is supplied.
+
+        Args:
+          depth (int): number of layers in the transformer
+          embedding_dim (int): the channel dimension for the input embeddings
+          num_heads (int): the number of heads for multihead attention. Must
+            divide embedding_dim
+          mlp_dim (int): the channel dimension internal to the MLP block
+          activation (nn.Module): the activation to use in the MLP block
+        """
+        super().__init__()
+        self.depth = depth
+        self.embedding_dim = embedding_dim
+        self.num_heads = num_heads
+        self.mlp_dim = mlp_dim
+        self.layers = nn.ModuleList()
+
+        for i in range(depth):
+            self.layers.append(
+                TwoWayAttentionBlock(
+                    embedding_dim=embedding_dim,
+                    num_heads=num_heads,
+                    mlp_dim=mlp_dim,
+                    activation=activation,
+                    attention_downsample_rate=attention_downsample_rate,
+                    skip_first_layer_pe=(i == 0),
+                )
+            )
+
+        self.final_attn_token_to_image = Attention(
+            embedding_dim, num_heads, downsample_rate=attention_downsample_rate
+        )
+        self.norm_final_attn = nn.LayerNorm(embedding_dim)
+
+    def forward(
+        self,
+        image_embedding: Tensor,
+        image_pe: Tensor,
+        point_embedding: Tensor,
+    ) -> Tuple[Tensor, Tensor]:
+        """
+        Args:
+          image_embedding (torch.Tensor): image to attend to. Should be shape
+            B x embedding_dim x h x w for any h and w.
+          image_pe (torch.Tensor): the positional encoding to add to the image. Must
+            have the same shape as image_embedding.
+          point_embedding (torch.Tensor): the embedding to add to the query points.
+            Must have shape B x N_points x embedding_dim for any N_points.
+
+        Returns:
+          torch.Tensor: the processed point_embedding
+          torch.Tensor: the processed image_embedding
+        """
+        # BxCxHxW -> BxHWxC == B x N_image_tokens x C
+        bs, c, h, w = image_embedding.shape
+        image_embedding = image_embedding.flatten(2).permute(0, 2, 1)
+        image_pe = image_pe.flatten(2).permute(0, 2, 1)
+
+        # Prepare queries
+        queries = point_embedding
+        keys = image_embedding
+
+        # Apply transformer blocks and final layernorm
+        for layer in self.layers:
+            queries, keys = layer(
+                queries=queries,
+                keys=keys,
+                query_pe=point_embedding,
+                key_pe=image_pe,
+            )
+
+        # Apply the final attention layer from the points to the image
+        q = queries + point_embedding
+        k = keys + image_pe
+        attn_out = self.final_attn_token_to_image(q=q, k=k, v=keys)
+        queries = queries + attn_out
+        queries = self.norm_final_attn(queries)
+
+        return queries, keys
+
+
+class TwoWayAttentionBlock(nn.Module):
+    def __init__(
+        self,
+        embedding_dim: int,
+        num_heads: int,
+        mlp_dim: int = 2048,
+        activation: Type[nn.Module] = nn.ReLU,
+        attention_downsample_rate: int = 2,
+        skip_first_layer_pe: bool = False,
+    ) -> None:
+        """
+        A transformer block with four layers: (1) self-attention of sparse
+        inputs, (2) cross attention of sparse inputs to dense inputs, (3) mlp
+        block on sparse inputs, and (4) cross attention of dense inputs to sparse
+        inputs.
+
+        Arguments:
+          embedding_dim (int): the channel dimension of the embeddings
+          num_heads (int): the number of heads in the attention layers
+          mlp_dim (int): the hidden dimension of the mlp block
+          activation (nn.Module): the activation of the mlp block
+          skip_first_layer_pe (bool): skip the PE on the first layer
+        """
+        super().__init__()
+        self.self_attn = Attention(embedding_dim, num_heads)
+        self.norm1 = nn.LayerNorm(embedding_dim)
+
+        self.cross_attn_token_to_image = Attention(
+            embedding_dim, num_heads, downsample_rate=attention_downsample_rate
+        )
+        self.norm2 = nn.LayerNorm(embedding_dim)
+
+        self.mlp = MLP(
+            embedding_dim, mlp_dim, embedding_dim, num_layers=2, activation=activation
+        )
+        self.norm3 = nn.LayerNorm(embedding_dim)
+
+        self.norm4 = nn.LayerNorm(embedding_dim)
+        self.cross_attn_image_to_token = Attention(
+            embedding_dim, num_heads, downsample_rate=attention_downsample_rate
+        )
+
+        self.skip_first_layer_pe = skip_first_layer_pe
+
+    def forward(
+        self, queries: Tensor, keys: Tensor, query_pe: Tensor, key_pe: Tensor
+    ) -> Tuple[Tensor, Tensor]:
+        # Self attention block
+        if self.skip_first_layer_pe:
+            queries = self.self_attn(q=queries, k=queries, v=queries)
+        else:
+            q = queries + query_pe
+            attn_out = self.self_attn(q=q, k=q, v=queries)
+            queries = queries + attn_out
+        queries = self.norm1(queries)
+
+        # Cross attention block, tokens attending to image embedding
+        q = queries + query_pe
+        k = keys + key_pe
+        attn_out = self.cross_attn_token_to_image(q=q, k=k, v=keys)
+        queries = queries + attn_out
+        queries = self.norm2(queries)
+
+        # MLP block
+        mlp_out = self.mlp(queries)
+        queries = queries + mlp_out
+        queries = self.norm3(queries)
+
+        # Cross attention block, image embedding attending to tokens
+        q = queries + query_pe
+        k = keys + key_pe
+        attn_out = self.cross_attn_image_to_token(q=k, k=q, v=queries)
+        keys = keys + attn_out
+        keys = self.norm4(keys)
+
+        return queries, keys
+
+
+class Attention(nn.Module):
+    """
+    An attention layer that allows for downscaling the size of the embedding
+    after projection to queries, keys, and values.
+    """
+
+    def __init__(
+        self,
+        embedding_dim: int,
+        num_heads: int,
+        downsample_rate: int = 1,
+        dropout: float = 0.0,
+        kv_in_dim: int = None,
+    ) -> None:
+        super().__init__()
+        self.embedding_dim = embedding_dim
+        self.kv_in_dim = kv_in_dim if kv_in_dim is not None else embedding_dim
+        self.internal_dim = embedding_dim // downsample_rate
+        self.num_heads = num_heads
+        assert (
+            self.internal_dim % num_heads == 0
+        ), "num_heads must divide embedding_dim."
+
+        self.q_proj = nn.Linear(embedding_dim, self.internal_dim)
+        self.k_proj = nn.Linear(self.kv_in_dim, self.internal_dim)
+        self.v_proj = nn.Linear(self.kv_in_dim, self.internal_dim)
+        self.out_proj = nn.Linear(self.internal_dim, embedding_dim)
+
+        self.dropout_p = dropout
+
+    def _separate_heads(self, x: Tensor, num_heads: int) -> Tensor:
+        b, n, c = x.shape
+        x = x.reshape(b, n, num_heads, c // num_heads)
+        return x.transpose(1, 2)  # B x N_heads x N_tokens x C_per_head
+
+    def _recombine_heads(self, x: Tensor) -> Tensor:
+        b, n_heads, n_tokens, c_per_head = x.shape
+        x = x.transpose(1, 2)
+        return x.reshape(b, n_tokens, n_heads * c_per_head)  # B x N_tokens x C
+
+    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
+        # Input projections
+        q = self.q_proj(q)
+        k = self.k_proj(k)
+        v = self.v_proj(v)
+
+        # Separate into heads
+        q = self._separate_heads(q, self.num_heads)
+        k = self._separate_heads(k, self.num_heads)
+        v = self._separate_heads(v, self.num_heads)
+
+        dropout_p = self.dropout_p if self.training else 0.0
+        if q.dtype != k.dtype: # these if statements added for mps support
+            k = k.to(dtype=q.dtype)
+        if q.dtype != v.dtype:
+            v = v.to(dtype=q.dtype)
+        # Attention
+        out = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p)
+
+        out = self._recombine_heads(out)
+        out = self.out_proj(out)
+
+        return out
+
+
+class RoPEAttention(Attention):
+    """Attention with rotary position encoding."""
+
+    def __init__(
+        self,
+        *args,
+        rope_theta=10000.0,
+        # whether to repeat q rope to match k length
+        # this is needed for cross-attention to memories
+        rope_k_repeat=False,
+        feat_sizes=(64, 64),  # [w, h] for stride 16 feats at 1024 resolution
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+
+        self.compute_cis = partial(
+            compute_axial_cis, dim=self.internal_dim // self.num_heads, theta=rope_theta
+        )
+        freqs_cis = self.compute_cis(end_x=feat_sizes[0], end_y=feat_sizes[1])
+        self.freqs_cis = (
+            freqs_cis.to("cuda") if torch.cuda.is_available() else freqs_cis
+        )
+        self.rope_k_repeat = rope_k_repeat
+
+    def forward(
+        self, q: Tensor, k: Tensor, v: Tensor, num_k_exclude_rope: int = 0
+    ) -> Tensor:
+        # Input projections
+        q = self.q_proj(q)
+        k = self.k_proj(k)
+        v = self.v_proj(v)
+
+        # Separate into heads
+        q = self._separate_heads(q, self.num_heads)
+        k = self._separate_heads(k, self.num_heads)
+        v = self._separate_heads(v, self.num_heads)
+
+        # Apply rotary position encoding
+        w = h = math.sqrt(q.shape[-2])
+        self.freqs_cis = self.freqs_cis.to(q.device)
+        if self.freqs_cis.shape[0] != q.shape[-2]:
+            self.freqs_cis = self.compute_cis(end_x=w, end_y=h).to(q.device)
+        if q.shape[-2] != k.shape[-2]:
+            assert self.rope_k_repeat
+
+        num_k_rope = k.size(-2) - num_k_exclude_rope
+        q, k[:, :, :num_k_rope] = apply_rotary_enc(
+            q,
+            k[:, :, :num_k_rope],
+            freqs_cis=self.freqs_cis,
+            repeat_freqs_k=self.rope_k_repeat,
+        )
+
+        dropout_p = self.dropout_p if self.training else 0.0
+        if q.dtype != k.dtype:
+            k = k.to(dtype=q.dtype)
+        if q.dtype != v.dtype:
+            v = v.to(dtype=q.dtype)
+        # Attention
+        out = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p)
+
+        out = self._recombine_heads(out)
+        out = self.out_proj(out)
+
+        return out
+
+class EfficientRoPEAttention1(Attention):
+    """Attention with rotary position encoding."""
+
+    def __init__(
+        self,
+        *args,
+        rope_theta=10000.0,
+        # whether to repeat q rope to match k length
+        # this is needed for cross-attention to memories
+        rope_k_repeat=False,
+        feat_sizes=(32, 32),  # [w, h] for stride 16 feats at 512 resolution
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+
+        self.compute_cis = partial(
+            compute_axial_cis, dim=self.internal_dim // self.num_heads, theta=rope_theta
+        )
+        freqs_cis = self.compute_cis(end_x=feat_sizes[0], end_y=feat_sizes[1])
+        self.freqs_cis = freqs_cis
+        self.rope_k_repeat = rope_k_repeat
+
+    def forward(
+        self, q: Tensor, k: Tensor, v: Tensor, num_k_exclude_rope: int = 0
+    ) -> Tensor:
+        # Input projections
+        q = self.q_proj(q)
+        k = self.k_proj(k)
+        v = self.v_proj(v)
+
+        # Separate into heads
+        q = self._separate_heads(q, self.num_heads)
+        k = self._separate_heads(k, self.num_heads)
+        v = self._separate_heads(v, self.num_heads)
+
+        # Apply rotary position encoding
+        w = h = math.sqrt(q.shape[-2])
+        self.freqs_cis = self.freqs_cis.to(q.device)
+        if self.freqs_cis.shape[0] != q.shape[-2]:
+            self.freqs_cis = self.compute_cis(end_x=w, end_y=h).to(q.device)
+        if q.shape[-2] != k.shape[-2]:
+            assert self.rope_k_repeat
+
+        num_k_rope = k.size(-2) - num_k_exclude_rope
+        q, k[:, :, :num_k_rope] = apply_rotary_enc(
+            q,
+            k[:, :, :num_k_rope],
+            freqs_cis=self.freqs_cis,
+            repeat_freqs_k=self.rope_k_repeat,
+        )
+
+        dropout_p = self.dropout_p if self.training else 0.0
+
+        if self.rope_k_repeat:
+            fs, bs, ns, ds = k.shape
+            nq = q.shape[-2]
+            if num_k_rope <= nq:
+                out = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p)
+            else:
+                s_kernel_size = 2
+                intw, inth = int(w), int(h)
+                k_landmarks = k[:, :, :num_k_rope, :].reshape(fs, -1, nq, ds)
+                k_landmarks = k_landmarks.transpose(-2, -1).reshape(fs, -1, intw, inth)
+                k_landmarks = F.avg_pool2d(
+                    k_landmarks, s_kernel_size, stride=s_kernel_size
+                )
+                k_landmarks = (
+                    k_landmarks.reshape(
+                        fs, -1, ds, nq // (s_kernel_size * s_kernel_size)
+                    )
+                    .transpose(-2, -1)
+                    .reshape(fs, bs, -1, ds)
+                )
+
+                scale_factor = 1 / math.sqrt(ds)
+                attn_weight = q @ k_landmarks.transpose(
+                    -2, -1
+                ) * scale_factor + 2 * math.log(s_kernel_size)
+                attn_weight = torch.cat(
+                    [
+                        attn_weight,
+                        q @ k[:, :, num_k_rope:, :].transpose(-2, -1) * scale_factor,
+                    ],
+                    dim=-1,
+                )
+                attn_weight = torch.softmax(attn_weight, dim=-1)
+                attn_weight = torch.dropout(attn_weight, dropout_p, train=self.training)
+
+                v_landmarks = v[:, :, :num_k_rope, :].reshape(fs, -1, nq, ds)
+                v_landmarks = v_landmarks.transpose(-2, -1).reshape(fs, -1, intw, inth)
+                v_landmarks = F.avg_pool2d(
+                    v_landmarks, s_kernel_size, stride=s_kernel_size
+                )
+                v_landmarks = v_landmarks.reshape(
+                    fs, -1, ds, nq // (s_kernel_size * s_kernel_size)
+                ).transpose(-2, -1)
+                v_landmarks = torch.cat(
+                    [
+                        v_landmarks.reshape(fs, bs, -1, ds),
+                        v[:, :, num_k_rope:, :],
+                    ],
+                    dim=-2,
+                )
+                out = attn_weight @ v_landmarks
+        else:
+            out = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p)
+
+        out = self._recombine_heads(out)
+        out = self.out_proj(out)
+
+        return out
+
+
+class EfficientRoPEAttention2(Attention):
+    """Attention with rotary position encoding."""
+
+    def __init__(
+        self,
+        *args,
+        rope_theta=10000.0,
+        # whether to repeat q rope to match k length
+        # this is needed for cross-attention to memories
+        rope_k_repeat=False,
+        feat_sizes=(32, 32),  # [w, h] for stride 16 feats at 512 resolution
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+
+        self.compute_cis = partial(
+            compute_axial_cis, dim=self.internal_dim // self.num_heads, theta=rope_theta
+        )
+        freqs_cis = self.compute_cis(end_x=feat_sizes[0], end_y=feat_sizes[1])
+        self.freqs_cis = freqs_cis
+        self.rope_k_repeat = rope_k_repeat
+
+    def forward(
+        self, q: Tensor, k: Tensor, v: Tensor, num_k_exclude_rope: int = 0
+    ) -> Tensor:
+        # Input projections
+        q = self.q_proj(q)
+        k = self.k_proj(k)
+        v = self.v_proj(v)
+
+        # Separate into heads
+        q = self._separate_heads(q, self.num_heads)
+        k = self._separate_heads(k, self.num_heads)
+        v = self._separate_heads(v, self.num_heads)
+
+        # Apply rotary position encoding
+        w = h = math.sqrt(q.shape[-2])
+        self.freqs_cis = self.freqs_cis.to(q.device)
+        if self.freqs_cis.shape[0] != q.shape[-2]:
+            self.freqs_cis = self.compute_cis(end_x=w, end_y=h).to(q.device)
+        if q.shape[-2] != k.shape[-2]:
+            assert self.rope_k_repeat
+
+        num_k_rope = k.size(-2) - num_k_exclude_rope
+        q, k[:, :, :num_k_rope] = apply_rotary_enc(
+            q,
+            k[:, :, :num_k_rope],
+            freqs_cis=self.freqs_cis,
+            repeat_freqs_k=self.rope_k_repeat,
+        )
+
+        dropout_p = self.dropout_p if self.training else 0.0
+
+        if self.rope_k_repeat:
+            fs, bs, ns, ds = k.shape
+            nq = q.shape[-2]
+            if num_k_rope <= nq:
+                out = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p)
+            else:
+                s_kernel_size = 2
+                intw, inth = int(w), int(h)
+                k_landmarks = k[:, :, :num_k_rope, :].reshape(fs, -1, nq, ds)
+                k_landmarks = k_landmarks.transpose(-2, -1).reshape(fs, -1, intw, inth)
+                k_landmarks = F.avg_pool2d(
+                    k_landmarks, s_kernel_size, stride=s_kernel_size
+                )
+                k_landmarks = k_landmarks.reshape(
+                    fs, -1, ds, nq // (s_kernel_size * s_kernel_size)
+                ).transpose(-2, -1)
+                k_landmarks = torch.cat(
+                    [
+                        k_landmarks.reshape(fs, bs, -1, ds)
+                        + 2 * math.log(s_kernel_size),
+                        k[:, :, num_k_rope:, :],
+                    ],
+                    dim=-2,
+                )
+
+                v_landmarks = v[:, :, :num_k_rope, :].reshape(fs, -1, nq, ds)
+                v_landmarks = v_landmarks.transpose(-2, -1).reshape(fs, -1, intw, inth)
+                v_landmarks = F.avg_pool2d(
+                    v_landmarks, s_kernel_size, stride=s_kernel_size
+                )
+                v_landmarks = v_landmarks.reshape(
+                    fs, -1, ds, nq // (s_kernel_size * s_kernel_size)
+                ).transpose(-2, -1)
+                v_landmarks = torch.cat(
+                    [
+                        v_landmarks.reshape(fs, bs, -1, ds),
+                        v[:, :, num_k_rope:, :],
+                    ],
+                    dim=-2,
+                )
+                out = F.scaled_dot_product_attention(
+                    q, k_landmarks, v_landmarks, dropout_p=dropout_p
+                )
+        else:
+            out = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p)
+
+        out = self._recombine_heads(out)
+        out = self.out_proj(out)
+
+        return out

utility/sam/modeling/sam2_base.py

@@ -0,0 +1,910 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import torch.distributed
+import torch.nn.functional as F
+
+from torch.nn.init import trunc_normal_
+
+from sam2.modeling.sam.mask_decoder import MaskDecoder
+from sam2.modeling.sam.prompt_encoder import PromptEncoder
+from sam2.modeling.sam.transformer import TwoWayTransformer
+from sam2.modeling.sam2_utils import get_1d_sine_pe, MLP, select_closest_cond_frames
+
+# a large negative value as a placeholder score for missing objects
+NO_OBJ_SCORE = -1024.0
+
+
+class SAM2Base(torch.nn.Module):
+    def __init__(
+        self,
+        image_encoder,
+        memory_attention,
+        memory_encoder,
+        num_maskmem=7,  # default 1 input frame + 6 previous frames
+        image_size=512,
+        backbone_stride=16,  # stride of the image backbone output
+        sigmoid_scale_for_mem_enc=1.0,  # scale factor for mask sigmoid prob
+        sigmoid_bias_for_mem_enc=0.0,  # bias factor for mask sigmoid prob
+        # During evaluation, whether to binarize the sigmoid mask logits on interacted frames with clicks
+        binarize_mask_from_pts_for_mem_enc=False,
+        use_mask_input_as_output_without_sam=False,  # on frames with mask input, whether to directly output the input mask without using a SAM prompt encoder + mask decoder
+        # The maximum number of conditioning frames to participate in the memory attention (-1 means no limit; if there are more conditioning frames than this limit,
+        # we only cross-attend to the temporally closest `max_cond_frames_in_attn` conditioning frames in the encoder when tracking each frame). This gives the model
+        # a temporal locality when handling a large number of annotated frames (since closer frames should be more important) and also avoids GPU OOM.
+        max_cond_frames_in_attn=-1,
+        # on the first frame, whether to directly add the no-memory embedding to the image feature
+        # (instead of using the transformer encoder)
+        directly_add_no_mem_embed=False,
+        # whether to use high-resolution feature maps in the SAM mask decoder
+        use_high_res_features_in_sam=False,
+        # whether to output multiple (3) masks for the first click on initial conditioning frames
+        multimask_output_in_sam=False,
+        # the minimum and maximum number of clicks to use multimask_output_in_sam (only relevant when `multimask_output_in_sam=True`;
+        # default is 1 for both, meaning that only the first click gives multimask output; also note that a box counts as two points)
+        multimask_min_pt_num=1,
+        multimask_max_pt_num=1,
+        # whether to also use multimask output for tracking (not just for the first click on initial conditioning frames; only relevant when `multimask_output_in_sam=True`)
+        multimask_output_for_tracking=False,
+        # Whether to use multimask tokens for obj ptr; Only relevant when both
+        # use_obj_ptrs_in_encoder=True and multimask_output_for_tracking=True
+        use_multimask_token_for_obj_ptr: bool = False,
+        # whether to use sigmoid to restrict ious prediction to [0-1]
+        iou_prediction_use_sigmoid=False,
+        # The memory bank's temporal stride during evaluation (i.e. the `r` parameter in XMem and Cutie; XMem and Cutie use r=5).
+        # For r>1, the (self.num_maskmem - 1) non-conditioning memory frames consist of
+        # (self.num_maskmem - 2) nearest frames from every r-th frames, plus the last frame.
+        memory_temporal_stride_for_eval=1,
+        # whether to apply non-overlapping constraints on the object masks in the memory encoder during evaluation (to avoid/alleviate superposing masks)
+        non_overlap_masks_for_mem_enc=False,
+        # whether to cross-attend to object pointers from other frames (based on SAM output tokens) in the encoder
+        use_obj_ptrs_in_encoder=False,
+        # the maximum number of object pointers from other frames in encoder cross attention (only relevant when `use_obj_ptrs_in_encoder=True`)
+        max_obj_ptrs_in_encoder=16,
+        # whether to add temporal positional encoding to the object pointers in the encoder (only relevant when `use_obj_ptrs_in_encoder=True`)
+        add_tpos_enc_to_obj_ptrs=True,
+        # whether to add an extra linear projection layer for the temporal positional encoding in the object pointers to avoid potential interference
+        # with spatial positional encoding (only relevant when both `use_obj_ptrs_in_encoder=True` and `add_tpos_enc_to_obj_ptrs=True`)
+        proj_tpos_enc_in_obj_ptrs=False,
+        # whether to use signed distance (instead of unsigned absolute distance) in the temporal positional encoding in the object pointers
+        # (only relevant when both `use_obj_ptrs_in_encoder=True` and `add_tpos_enc_to_obj_ptrs=True`)
+        use_signed_tpos_enc_to_obj_ptrs=False,
+        # whether to only attend to object pointers in the past (before the current frame) in the encoder during evaluation
+        # (only relevant when `use_obj_ptrs_in_encoder=True`; this might avoid pointer information too far in the future to distract the initial tracking)
+        only_obj_ptrs_in_the_past_for_eval=False,
+        # Whether to predict if there is an object in the frame
+        pred_obj_scores: bool = False,
+        # Whether to use an MLP to predict object scores
+        pred_obj_scores_mlp: bool = False,
+        # Only relevant if pred_obj_scores=True and use_obj_ptrs_in_encoder=True;
+        # Whether to have a fixed no obj pointer when there is no object present
+        # or to use it as an additive embedding with obj_ptr produced by decoder
+        fixed_no_obj_ptr: bool = False,
+        # Soft no object, i.e. mix in no_obj_ptr softly,
+        # hope to make recovery easier if there is a mistake and mitigate accumulation of errors
+        soft_no_obj_ptr: bool = False,
+        use_mlp_for_obj_ptr_proj: bool = False,
+        # add no obj embedding to spatial frames
+        no_obj_embed_spatial: bool = False,
+        # extra arguments used to construct the SAM mask decoder; if not None, it should be a dict of kwargs to be passed into `MaskDecoder` class.
+        sam_mask_decoder_extra_args=None,
+        compile_image_encoder: bool = False,
+    ):
+        super().__init__()
+
+        # Part 1: the image backbone
+        self.image_encoder = image_encoder
+        # Use level 0, 1, 2 for high-res setting, or just level 2 for the default setting
+        self.use_high_res_features_in_sam = use_high_res_features_in_sam
+        self.num_feature_levels = 3 if use_high_res_features_in_sam else 1
+        self.use_obj_ptrs_in_encoder = use_obj_ptrs_in_encoder
+        self.max_obj_ptrs_in_encoder = max_obj_ptrs_in_encoder
+        if use_obj_ptrs_in_encoder:
+            # A conv layer to downsample the mask prompt to stride 4 (the same stride as
+            # low-res SAM mask logits) and to change its scales from 0~1 to SAM logit scale,
+            # so that it can be fed into the SAM mask decoder to generate a pointer.
+            self.mask_downsample = torch.nn.Conv2d(1, 1, kernel_size=4, stride=4)
+        self.add_tpos_enc_to_obj_ptrs = add_tpos_enc_to_obj_ptrs
+        if proj_tpos_enc_in_obj_ptrs:
+            assert add_tpos_enc_to_obj_ptrs  # these options need to be used together
+        self.proj_tpos_enc_in_obj_ptrs = proj_tpos_enc_in_obj_ptrs
+        self.use_signed_tpos_enc_to_obj_ptrs = use_signed_tpos_enc_to_obj_ptrs
+        self.only_obj_ptrs_in_the_past_for_eval = only_obj_ptrs_in_the_past_for_eval
+
+        # Part 2: memory attention to condition current frame's visual features
+        # with memories (and obj ptrs) from past frames
+        self.memory_attention = memory_attention
+        self.hidden_dim = image_encoder.neck.d_model
+
+        # Part 3: memory encoder for the previous frame's outputs
+        self.memory_encoder = memory_encoder
+        self.mem_dim = self.hidden_dim
+        if hasattr(self.memory_encoder, "out_proj") and hasattr(
+            self.memory_encoder.out_proj, "weight"
+        ):
+            # if there is compression of memories along channel dim
+            self.mem_dim = self.memory_encoder.out_proj.weight.shape[0]
+        self.num_maskmem = num_maskmem  # Number of memories accessible
+        # Temporal encoding of the memories
+        self.maskmem_tpos_enc = torch.nn.Parameter(
+            torch.zeros(num_maskmem, 1, 1, self.mem_dim)
+        )
+        trunc_normal_(self.maskmem_tpos_enc, std=0.02)
+        # a single token to indicate no memory embedding from previous frames
+        self.no_mem_embed = torch.nn.Parameter(torch.zeros(1, 1, self.hidden_dim))
+        self.no_mem_pos_enc = torch.nn.Parameter(torch.zeros(1, 1, self.hidden_dim))
+        trunc_normal_(self.no_mem_embed, std=0.02)
+        trunc_normal_(self.no_mem_pos_enc, std=0.02)
+        self.directly_add_no_mem_embed = directly_add_no_mem_embed
+        # Apply sigmoid to the output raw mask logits (to turn them from
+        # range (-inf, +inf) to range (0, 1)) before feeding them into the memory encoder
+        self.sigmoid_scale_for_mem_enc = sigmoid_scale_for_mem_enc
+        self.sigmoid_bias_for_mem_enc = sigmoid_bias_for_mem_enc
+        self.binarize_mask_from_pts_for_mem_enc = binarize_mask_from_pts_for_mem_enc
+        self.non_overlap_masks_for_mem_enc = non_overlap_masks_for_mem_enc
+        self.memory_temporal_stride_for_eval = memory_temporal_stride_for_eval
+        # On frames with mask input, whether to directly output the input mask without
+        # using a SAM prompt encoder + mask decoder
+        self.use_mask_input_as_output_without_sam = use_mask_input_as_output_without_sam
+        self.multimask_output_in_sam = multimask_output_in_sam
+        self.multimask_min_pt_num = multimask_min_pt_num
+        self.multimask_max_pt_num = multimask_max_pt_num
+        self.multimask_output_for_tracking = multimask_output_for_tracking
+        self.use_multimask_token_for_obj_ptr = use_multimask_token_for_obj_ptr
+        self.iou_prediction_use_sigmoid = iou_prediction_use_sigmoid
+
+        # Part 4: SAM-style prompt encoder (for both mask and point inputs)
+        # and SAM-style mask decoder for the final mask output
+        self.image_size = image_size
+        self.backbone_stride = backbone_stride
+        self.sam_mask_decoder_extra_args = sam_mask_decoder_extra_args
+        self.pred_obj_scores = pred_obj_scores
+        self.pred_obj_scores_mlp = pred_obj_scores_mlp
+        self.fixed_no_obj_ptr = fixed_no_obj_ptr
+        self.soft_no_obj_ptr = soft_no_obj_ptr
+        if self.fixed_no_obj_ptr:
+            assert self.pred_obj_scores
+            assert self.use_obj_ptrs_in_encoder
+        if self.pred_obj_scores and self.use_obj_ptrs_in_encoder:
+            self.no_obj_ptr = torch.nn.Parameter(torch.zeros(1, self.hidden_dim))
+            trunc_normal_(self.no_obj_ptr, std=0.02)
+        self.use_mlp_for_obj_ptr_proj = use_mlp_for_obj_ptr_proj
+        self.no_obj_embed_spatial = None
+        if no_obj_embed_spatial:
+            self.no_obj_embed_spatial = torch.nn.Parameter(torch.zeros(1, self.mem_dim))
+            trunc_normal_(self.no_obj_embed_spatial, std=0.02)
+
+        self._build_sam_heads()
+        self.max_cond_frames_in_attn = max_cond_frames_in_attn
+
+        # Model compilation
+        if compile_image_encoder:
+            # Compile the forward function (not the full module) to allow loading checkpoints.
+            print(
+                "Image encoder compilation is enabled. First forward pass will be slow."
+            )
+            self.image_encoder.forward = torch.compile(
+                self.image_encoder.forward,
+                mode="max-autotune",
+                fullgraph=True,
+                dynamic=False,
+            )
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    def forward(self, *args, **kwargs):
+        raise NotImplementedError(
+            "Please use the corresponding methods in SAM2VideoPredictor for inference or SAM2Train for training/fine-tuning"
+            "See notebooks/video_predictor_example.ipynb for an inference example."
+        )
+
+    def _build_sam_heads(self):
+        """Build SAM-style prompt encoder and mask decoder."""
+        self.sam_prompt_embed_dim = self.hidden_dim
+        self.sam_image_embedding_size = self.image_size // self.backbone_stride
+
+        # build PromptEncoder and MaskDecoder from SAM
+        # (their hyperparameters like `mask_in_chans=16` are from SAM code)
+        self.sam_prompt_encoder = PromptEncoder(
+            embed_dim=self.sam_prompt_embed_dim,
+            image_embedding_size=(
+                self.sam_image_embedding_size,
+                self.sam_image_embedding_size,
+            ),
+            input_image_size=(self.image_size, self.image_size),
+            mask_in_chans=16,
+        )
+        self.sam_mask_decoder = MaskDecoder(
+            num_multimask_outputs=3,
+            transformer=TwoWayTransformer(
+                depth=2,
+                embedding_dim=self.sam_prompt_embed_dim,
+                mlp_dim=2048,
+                num_heads=8,
+            ),
+            transformer_dim=self.sam_prompt_embed_dim,
+            iou_head_depth=3,
+            iou_head_hidden_dim=256,
+            use_high_res_features=self.use_high_res_features_in_sam,
+            iou_prediction_use_sigmoid=self.iou_prediction_use_sigmoid,
+            pred_obj_scores=self.pred_obj_scores,
+            pred_obj_scores_mlp=self.pred_obj_scores_mlp,
+            use_multimask_token_for_obj_ptr=self.use_multimask_token_for_obj_ptr,
+            **(self.sam_mask_decoder_extra_args or {}),
+        )
+        if self.use_obj_ptrs_in_encoder:
+            # a linear projection on SAM output tokens to turn them into object pointers
+            self.obj_ptr_proj = torch.nn.Linear(self.hidden_dim, self.hidden_dim)
+            if self.use_mlp_for_obj_ptr_proj:
+                self.obj_ptr_proj = MLP(
+                    self.hidden_dim, self.hidden_dim, self.hidden_dim, 3
+                )
+        else:
+            self.obj_ptr_proj = torch.nn.Identity()
+        if self.proj_tpos_enc_in_obj_ptrs:
+            # a linear projection on temporal positional encoding in object pointers to
+            # avoid potential interference with spatial positional encoding
+            self.obj_ptr_tpos_proj = torch.nn.Linear(self.hidden_dim, self.mem_dim)
+        else:
+            self.obj_ptr_tpos_proj = torch.nn.Identity()
+
+    def _forward_sam_heads(
+        self,
+        backbone_features,
+        point_inputs=None,
+        mask_inputs=None,
+        high_res_features=None,
+        multimask_output=False,
+    ):
+        """
+        Forward SAM prompt encoders and mask heads.
+
+        Inputs:
+        - backbone_features: image features of [B, C, H, W] shape
+        - point_inputs: a dictionary with "point_coords" and "point_labels", where
+          1) "point_coords" has [B, P, 2] shape and float32 dtype and contains the
+             absolute pixel-unit coordinate in (x, y) format of the P input points
+          2) "point_labels" has shape [B, P] and int32 dtype, where 1 means
+             positive clicks, 0 means negative clicks, and -1 means padding
+        - mask_inputs: a mask of [B, 1, H*16, W*16] shape, float or bool, with the
+          same spatial size as the image.
+        - high_res_features: either 1) None or 2) or a list of length 2 containing
+          two feature maps of [B, C, 4*H, 4*W] and [B, C, 2*H, 2*W] shapes respectively,
+          which will be used as high-resolution feature maps for SAM decoder.
+        - multimask_output: if it's True, we output 3 candidate masks and their 3
+          corresponding IoU estimates, and if it's False, we output only 1 mask and
+          its corresponding IoU estimate.
+
+        Outputs:
+        - low_res_multimasks: [B, M, H*4, W*4] shape (where M = 3 if
+          `multimask_output=True` and M = 1 if `multimask_output=False`), the SAM
+          output mask logits (before sigmoid) for the low-resolution masks, with 4x
+          the resolution (1/4 stride) of the input backbone_features.
+        - high_res_multimasks: [B, M, H*16, W*16] shape (where M = 3
+          if `multimask_output=True` and M = 1 if `multimask_output=False`),
+          upsampled from the low-resolution masks, with shape size as the image
+          (stride is 1 pixel).
+        - ious, [B, M] shape, where (where M = 3 if `multimask_output=True` and M = 1
+          if `multimask_output=False`), the estimated IoU of each output mask.
+        - low_res_masks: [B, 1, H*4, W*4] shape, the best mask in `low_res_multimasks`.
+          If `multimask_output=True`, it's the mask with the highest IoU estimate.
+          If `multimask_output=False`, it's the same as `low_res_multimasks`.
+        - high_res_masks: [B, 1, H*16, W*16] shape, the best mask in `high_res_multimasks`.
+          If `multimask_output=True`, it's the mask with the highest IoU estimate.
+          If `multimask_output=False`, it's the same as `high_res_multimasks`.
+        - obj_ptr: [B, C] shape, the object pointer vector for the output mask, extracted
+          based on the output token from the SAM mask decoder.
+        """
+        B = backbone_features.size(0)
+        device = backbone_features.device
+        assert backbone_features.size(1) == self.sam_prompt_embed_dim
+        assert backbone_features.size(2) == self.sam_image_embedding_size
+        assert backbone_features.size(3) == self.sam_image_embedding_size
+
+        # a) Handle point prompts
+        if point_inputs is not None:
+            sam_point_coords = point_inputs["point_coords"]
+            sam_point_labels = point_inputs["point_labels"]
+            assert sam_point_coords.size(0) == B and sam_point_labels.size(0) == B
+        else:
+            # If no points are provide, pad with an empty point (with label -1)
+            sam_point_coords = torch.zeros(B, 1, 2, device=device)
+            sam_point_labels = -torch.ones(B, 1, dtype=torch.int32, device=device)
+
+        # b) Handle mask prompts
+        if mask_inputs is not None:
+            # If mask_inputs is provided, downsize it into low-res mask input if needed
+            # and feed it as a dense mask prompt into the SAM mask encoder
+            assert len(mask_inputs.shape) == 4 and mask_inputs.shape[:2] == (B, 1)
+            if mask_inputs.shape[-2:] != self.sam_prompt_encoder.mask_input_size:
+                sam_mask_prompt = F.interpolate(
+                    mask_inputs.float(),
+                    size=self.sam_prompt_encoder.mask_input_size,
+                    align_corners=False,
+                    mode="bilinear",
+                    antialias=True,  # use antialias for downsampling
+                )
+            else:
+                sam_mask_prompt = mask_inputs
+        else:
+            # Otherwise, simply feed None (and SAM's prompt encoder will add
+            # a learned `no_mask_embed` to indicate no mask input in this case).
+            sam_mask_prompt = None
+
+        sparse_embeddings, dense_embeddings = self.sam_prompt_encoder(
+            points=(sam_point_coords, sam_point_labels),
+            boxes=None,
+            masks=sam_mask_prompt,
+        )
+        (
+            low_res_multimasks,
+            ious,
+            sam_output_tokens,
+            object_score_logits,
+        ) = self.sam_mask_decoder(
+            image_embeddings=backbone_features,
+            image_pe=self.sam_prompt_encoder.get_dense_pe(),
+            sparse_prompt_embeddings=sparse_embeddings,
+            dense_prompt_embeddings=dense_embeddings,
+            multimask_output=multimask_output,
+            repeat_image=False,  # the image is already batched
+            high_res_features=high_res_features,
+        )
+        if self.pred_obj_scores:
+            is_obj_appearing = object_score_logits > 0
+
+            # Mask used for spatial memories is always a *hard* choice between obj and no obj,
+            # consistent with the actual mask prediction
+            low_res_multimasks = torch.where(
+                is_obj_appearing[:, None, None],
+                low_res_multimasks,
+                NO_OBJ_SCORE,
+            )
+
+        # convert masks from possibly bfloat16 (or float16) to float32
+        # (older PyTorch versions before 2.1 don't support `interpolate` on bf16)
+        low_res_multimasks = low_res_multimasks.float()
+        high_res_multimasks = F.interpolate(
+            low_res_multimasks,
+            size=(self.image_size, self.image_size),
+            mode="bilinear",
+            align_corners=False,
+        )
+
+        sam_output_token = sam_output_tokens[:, 0]
+        if multimask_output:
+            # take the best mask prediction (with the highest IoU estimation)
+            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
+
+        # Extract object pointer from the SAM output token (with occlusion handling)
+        obj_ptr = self.obj_ptr_proj(sam_output_token)
+        if self.pred_obj_scores:
+            # Allow *soft* no obj ptr, unlike for masks
+            if self.soft_no_obj_ptr:
+                lambda_is_obj_appearing = object_score_logits.sigmoid()
+            else:
+                lambda_is_obj_appearing = is_obj_appearing.float()
+
+            if self.fixed_no_obj_ptr:
+                obj_ptr = lambda_is_obj_appearing * obj_ptr
+            obj_ptr = obj_ptr + (1 - lambda_is_obj_appearing) * self.no_obj_ptr
+
+        return (
+            low_res_multimasks,
+            high_res_multimasks,
+            ious,
+            low_res_masks,
+            high_res_masks,
+            obj_ptr,
+            object_score_logits,
+        )
+
+    def _use_mask_as_output(self, backbone_features, high_res_features, mask_inputs):
+        """
+        Directly turn binary `mask_inputs` into a output mask logits without using SAM.
+        (same input and output shapes as in _forward_sam_heads above).
+        """
+        # Use -10/+10 as logits for neg/pos pixels (very close to 0/1 in prob after sigmoid).
+        out_scale, out_bias = 20.0, -10.0  # sigmoid(-10.0)=4.5398e-05
+        mask_inputs_float = mask_inputs.float()
+        high_res_masks = mask_inputs_float * out_scale + out_bias
+        low_res_masks = F.interpolate(
+            high_res_masks,
+            size=(high_res_masks.size(-2) // 4, high_res_masks.size(-1) // 4),
+            align_corners=False,
+            mode="bilinear",
+            antialias=True,  # use antialias for downsampling
+        )
+        # a dummy IoU prediction of all 1's under mask input
+        ious = mask_inputs.new_ones(mask_inputs.size(0), 1).float()
+        if not self.use_obj_ptrs_in_encoder:
+            # all zeros as a dummy object pointer (of shape [B, C])
+            obj_ptr = torch.zeros(
+                mask_inputs.size(0), self.hidden_dim, device=mask_inputs.device
+            )
+        else:
+            # produce an object pointer using the SAM decoder from the mask input
+            _, _, _, _, _, obj_ptr, _ = self._forward_sam_heads(
+                backbone_features=backbone_features,
+                mask_inputs=self.mask_downsample(mask_inputs_float),
+                high_res_features=high_res_features,
+            )
+        # In this method, we are treating mask_input as output, e.g. using it directly to create spatial mem;
+        # Below, we follow the same design axiom to use mask_input to decide if obj appears or not instead of relying
+        # on the object_scores from the SAM decoder.
+        is_obj_appearing = torch.any(mask_inputs.flatten(1).float() > 0.0, dim=1)
+        is_obj_appearing = is_obj_appearing[..., None]
+        lambda_is_obj_appearing = is_obj_appearing.float()
+        object_score_logits = out_scale * lambda_is_obj_appearing + out_bias
+        if self.pred_obj_scores:
+            if self.fixed_no_obj_ptr:
+                obj_ptr = lambda_is_obj_appearing * obj_ptr
+            obj_ptr = obj_ptr + (1 - lambda_is_obj_appearing) * self.no_obj_ptr
+
+        return (
+            low_res_masks,
+            high_res_masks,
+            ious,
+            low_res_masks,
+            high_res_masks,
+            obj_ptr,
+            object_score_logits,
+        )
+
+    def forward_image(self, img_batch: torch.Tensor):
+        """Get the image feature on the input batch."""
+        backbone_out = self.image_encoder(img_batch)
+        if self.use_high_res_features_in_sam:
+            # precompute projected level 0 and level 1 features in SAM decoder
+            # to avoid running it again on every SAM click
+            backbone_out["backbone_fpn"][0] = self.sam_mask_decoder.conv_s0(
+                backbone_out["backbone_fpn"][0]
+            )
+            backbone_out["backbone_fpn"][1] = self.sam_mask_decoder.conv_s1(
+                backbone_out["backbone_fpn"][1]
+            )
+        return backbone_out
+
+    def _prepare_backbone_features(self, backbone_out):
+        """Prepare and flatten visual features."""
+        backbone_out = backbone_out.copy()
+        assert len(backbone_out["backbone_fpn"]) == len(backbone_out["vision_pos_enc"])
+        assert len(backbone_out["backbone_fpn"]) >= self.num_feature_levels
+
+        feature_maps = backbone_out["backbone_fpn"][-self.num_feature_levels :]
+        vision_pos_embeds = backbone_out["vision_pos_enc"][-self.num_feature_levels :]
+
+        feat_sizes = [(x.shape[-2], x.shape[-1]) for x in vision_pos_embeds]
+        # flatten NxCxHxW to HWxNxC
+        vision_feats = [x.flatten(2).permute(2, 0, 1) for x in feature_maps]
+        vision_pos_embeds = [x.flatten(2).permute(2, 0, 1) for x in vision_pos_embeds]
+
+        return backbone_out, vision_feats, vision_pos_embeds, feat_sizes
+
+    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,
+        track_in_reverse=False,  # tracking in reverse time order (for demo usage)
+    ):
+        """Fuse the current frame's visual feature map with previous memory."""
+        B = current_vision_feats[-1].size(1)  # batch size on this frame
+        C = self.hidden_dim
+        H, W = feat_sizes[-1]  # top-level (lowest-resolution) feature size
+        device = current_vision_feats[-1].device
+        # The case of `self.num_maskmem == 0` below is primarily used for reproducing SAM on images.
+        # In this case, we skip the fusion with any memory.
+        if self.num_maskmem == 0:  # Disable memory and skip fusion
+            pix_feat = current_vision_feats[-1].permute(1, 2, 0).view(B, C, H, W)
+            return pix_feat
+
+        num_obj_ptr_tokens = 0
+        tpos_sign_mul = -1 if track_in_reverse else 1
+        # Step 1: condition the visual features of the current frame on previous memories
+        if not is_init_cond_frame:
+            # Retrieve the memories encoded with the maskmem backbone
+            to_cat_memory, to_cat_memory_pos_embed = [], []
+            # Add conditioning frames's output first (all cond frames have t_pos=0 for
+            # when getting temporal positional embedding below)
+            assert len(output_dict["cond_frame_outputs"]) > 0
+            # Select a maximum number of temporally closest cond frames for cross attention
+            cond_outputs = output_dict["cond_frame_outputs"]
+            selected_cond_outputs, unselected_cond_outputs = select_closest_cond_frames(
+                frame_idx, cond_outputs, self.max_cond_frames_in_attn
+            )
+            t_pos_and_prevs = [(0, out) for out in selected_cond_outputs.values()]
+            # Add last (self.num_maskmem - 1) frames before current frame for non-conditioning memory
+            # the earliest one has t_pos=1 and the latest one has t_pos=self.num_maskmem-1
+            # We also allow taking the memory frame non-consecutively (with stride>1), in which case
+            # we take (self.num_maskmem - 2) frames among every stride-th frames plus the last frame.
+            stride = 1 if self.training else self.memory_temporal_stride_for_eval
+            for t_pos in range(1, self.num_maskmem):
+                t_rel = self.num_maskmem - t_pos  # how many frames before current frame
+                if t_rel == 1:
+                    # for t_rel == 1, we take the last frame (regardless of r)
+                    if not track_in_reverse:
+                        # the frame immediately before this frame (i.e. frame_idx - 1)
+                        prev_frame_idx = frame_idx - t_rel
+                    else:
+                        # the frame immediately after this frame (i.e. frame_idx + 1)
+                        prev_frame_idx = frame_idx + t_rel
+                else:
+                    # for t_rel >= 2, we take the memory frame from every r-th frames
+                    if not track_in_reverse:
+                        # first find the nearest frame among every r-th frames before this frame
+                        # for r=1, this would be (frame_idx - 2)
+                        prev_frame_idx = ((frame_idx - 2) // stride) * stride
+                        # then seek further among every r-th frames
+                        prev_frame_idx = prev_frame_idx - (t_rel - 2) * stride
+                    else:
+                        # first find the nearest frame among every r-th frames after this frame
+                        # for r=1, this would be (frame_idx + 2)
+                        prev_frame_idx = -(-(frame_idx + 2) // stride) * stride
+                        # then seek further among every r-th frames
+                        prev_frame_idx = prev_frame_idx + (t_rel - 2) * stride
+                out = output_dict["non_cond_frame_outputs"].get(prev_frame_idx, None)
+                if out is None:
+                    # If an unselected conditioning frame is among the last (self.num_maskmem - 1)
+                    # frames, we still attend to it as if it's a non-conditioning frame.
+                    out = unselected_cond_outputs.get(prev_frame_idx, None)
+                t_pos_and_prevs.append((t_pos, out))
+
+            for t_pos, prev in t_pos_and_prevs:
+                if prev is None:
+                    continue  # skip padding frames
+                # "maskmem_features" might have been offloaded to CPU in demo use cases,
+                # so we load it back to GPU (it's a no-op if it's already on GPU).
+                feats = prev["maskmem_features"].to(device, non_blocking=True)
+                feats = feats.to(torch.float32)
+                to_cat_memory.append(feats.flatten(2).permute(2, 0, 1))
+                # Spatial positional encoding (it might have been offloaded to CPU in eval)
+                maskmem_enc = prev["maskmem_pos_enc"][-1].to(device)
+                maskmem_enc = maskmem_enc.flatten(2).permute(2, 0, 1)
+                # Temporal positional encoding
+                maskmem_enc = (
+                    maskmem_enc + self.maskmem_tpos_enc[self.num_maskmem - t_pos - 1]
+                )
+                to_cat_memory_pos_embed.append(maskmem_enc)
+
+            # Construct the list of past object pointers
+            if self.use_obj_ptrs_in_encoder:
+                max_obj_ptrs_in_encoder = min(num_frames, self.max_obj_ptrs_in_encoder)
+                # First add those object pointers from selected conditioning frames
+                # (optionally, only include object pointers in the past during evaluation)
+                if not self.training and self.only_obj_ptrs_in_the_past_for_eval:
+                    ptr_cond_outputs = {
+                        t: out
+                        for t, out in selected_cond_outputs.items()
+                        if (t >= frame_idx if track_in_reverse else t <= frame_idx)
+                    }
+                else:
+                    ptr_cond_outputs = selected_cond_outputs
+                pos_and_ptrs = [
+                    # Temporal pos encoding contains how far away each pointer is from current frame
+                    (
+                        (
+                            (frame_idx - t) * tpos_sign_mul
+                            if self.use_signed_tpos_enc_to_obj_ptrs
+                            else abs(frame_idx - t)
+                        ),
+                        out["obj_ptr"],
+                    )
+                    for t, out in ptr_cond_outputs.items()
+                ]
+                # Add up to (max_obj_ptrs_in_encoder - 1) non-conditioning frames before current frame
+                for t_diff in range(1, max_obj_ptrs_in_encoder):
+                    t = frame_idx + t_diff if track_in_reverse else frame_idx - t_diff
+                    if t < 0 or (num_frames is not None and t >= num_frames):
+                        break
+                    out = output_dict["non_cond_frame_outputs"].get(
+                        t, unselected_cond_outputs.get(t, None)
+                    )
+                    if out is not None:
+                        pos_and_ptrs.append((t_diff, out["obj_ptr"]))
+                # If we have at least one object pointer, add them to the across attention
+                if len(pos_and_ptrs) > 0:
+                    pos_list, ptrs_list = zip(*pos_and_ptrs)
+                    # stack object pointers along dim=0 into [ptr_seq_len, B, C] shape
+                    obj_ptrs = torch.stack(ptrs_list, dim=0)
+                    # a temporal positional embedding based on how far each object pointer is from
+                    # the current frame (sine embedding normalized by the max pointer num).
+                    if self.add_tpos_enc_to_obj_ptrs:
+                        t_diff_max = max_obj_ptrs_in_encoder - 1
+                        tpos_dim = C if self.proj_tpos_enc_in_obj_ptrs else self.mem_dim
+                        obj_pos = torch.tensor(pos_list).to(
+                            device=device, non_blocking=True
+                        )
+                        obj_pos = get_1d_sine_pe(obj_pos / t_diff_max, dim=tpos_dim)
+                        obj_pos = self.obj_ptr_tpos_proj(obj_pos)
+                        obj_pos = obj_pos.unsqueeze(1).expand(-1, B, self.mem_dim)
+                    else:
+                        obj_pos = obj_ptrs.new_zeros(len(pos_list), B, self.mem_dim)
+                    if self.mem_dim < C:
+                        # split a pointer into (C // self.mem_dim) tokens for self.mem_dim < C
+                        obj_ptrs = obj_ptrs.reshape(
+                            -1, B, C // self.mem_dim, self.mem_dim
+                        )
+                        obj_ptrs = obj_ptrs.permute(0, 2, 1, 3).flatten(0, 1)
+                        obj_pos = obj_pos.repeat_interleave(C // self.mem_dim, dim=0)
+                    to_cat_memory.append(obj_ptrs)
+                    to_cat_memory_pos_embed.append(obj_pos)
+                    num_obj_ptr_tokens = obj_ptrs.shape[0]
+                else:
+                    num_obj_ptr_tokens = 0
+        else:
+            # for initial conditioning frames, encode them without using any previous memory
+            if self.directly_add_no_mem_embed:
+                # directly add no-mem embedding (instead of using the transformer encoder)
+                pix_feat_with_mem = current_vision_feats[-1] + self.no_mem_embed
+                pix_feat_with_mem = pix_feat_with_mem.permute(1, 2, 0).view(B, C, H, W)
+                return pix_feat_with_mem
+
+            # Use a dummy token on the first frame (to avoid empty memory input to tranformer encoder)
+            to_cat_memory = [self.no_mem_embed.expand(1, B, self.mem_dim)]
+            to_cat_memory_pos_embed = [self.no_mem_pos_enc.expand(1, B, self.mem_dim)]
+
+        # Step 2: Concatenate the memories and forward through the transformer encoder
+        memory = torch.cat(to_cat_memory, dim=0)
+        memory_pos_embed = torch.cat(to_cat_memory_pos_embed, dim=0)
+
+        pix_feat_with_mem = self.memory_attention(
+            curr=current_vision_feats,
+            curr_pos=current_vision_pos_embeds,
+            memory=memory,
+            memory_pos=memory_pos_embed,
+            num_obj_ptr_tokens=num_obj_ptr_tokens,
+        )
+        # reshape the output (HW)BC => BCHW
+        pix_feat_with_mem = pix_feat_with_mem.permute(1, 2, 0).view(B, C, H, W)
+        return pix_feat_with_mem
+
+    def _encode_new_memory(
+        self,
+        current_vision_feats,
+        feat_sizes,
+        pred_masks_high_res,
+        object_score_logits,
+        is_mask_from_pts,
+    ):
+        """Encode the current image and its prediction into a memory feature."""
+        B = current_vision_feats[-1].size(1)  # batch size on this frame
+        C = self.hidden_dim
+        H, W = feat_sizes[-1]  # top-level (lowest-resolution) feature size
+        # top-level feature, (HW)BC => BCHW
+        pix_feat = current_vision_feats[-1].permute(1, 2, 0).view(B, C, H, W)
+        if self.non_overlap_masks_for_mem_enc and not self.training:
+            # optionally, apply non-overlapping constraints to the masks (it's applied
+            # in the batch dimension and should only be used during eval, where all
+            # the objects come from the same video under batch size 1).
+            pred_masks_high_res = self._apply_non_overlapping_constraints(
+                pred_masks_high_res
+            )
+        # scale the raw mask logits with a temperature before applying sigmoid
+        binarize = self.binarize_mask_from_pts_for_mem_enc and is_mask_from_pts
+        if binarize and not self.training:
+            mask_for_mem = (pred_masks_high_res > 0).float()
+        else:
+            # apply sigmoid on the raw mask logits to turn them into range (0, 1)
+            mask_for_mem = torch.sigmoid(pred_masks_high_res)
+        # apply scale and bias terms to the sigmoid probabilities
+        if self.sigmoid_scale_for_mem_enc != 1.0:
+            mask_for_mem = mask_for_mem * self.sigmoid_scale_for_mem_enc
+        if self.sigmoid_bias_for_mem_enc != 0.0:
+            mask_for_mem = mask_for_mem + self.sigmoid_bias_for_mem_enc
+        maskmem_out = self.memory_encoder(
+            pix_feat, mask_for_mem, skip_mask_sigmoid=True  # sigmoid already applied
+        )
+        maskmem_features = maskmem_out["vision_features"]
+        maskmem_pos_enc = maskmem_out["vision_pos_enc"]
+        # add a no-object embedding to the spatial memory to indicate that the frame
+        # is predicted to be occluded (i.e. no object is appearing in the frame)
+        if self.no_obj_embed_spatial is not None:
+            is_obj_appearing = (object_score_logits > 0).float()
+            maskmem_features += (
+                1 - is_obj_appearing[..., None, None]
+            ) * self.no_obj_embed_spatial[..., None, None].expand(
+                *maskmem_features.shape
+            )
+
+        return maskmem_features, maskmem_pos_enc
+
+    def _track_step(
+        self,
+        frame_idx,
+        is_init_cond_frame,
+        current_vision_feats,
+        current_vision_pos_embeds,
+        feat_sizes,
+        point_inputs,
+        mask_inputs,
+        output_dict,
+        num_frames,
+        track_in_reverse,
+        prev_sam_mask_logits,
+    ):
+        current_out = {"point_inputs": point_inputs, "mask_inputs": mask_inputs}
+        # High-resolution feature maps for the SAM head, reshape (HW)BC => BCHW
+        if len(current_vision_feats) > 1:
+            high_res_features = [
+                x.permute(1, 2, 0).view(x.size(1), x.size(2), *s)
+                for x, s in zip(current_vision_feats[:-1], feat_sizes[:-1])
+            ]
+        else:
+            high_res_features = None
+        if mask_inputs is not None and self.use_mask_input_as_output_without_sam:
+            # When use_mask_input_as_output_without_sam=True, we directly output the mask input
+            # (see it as a GT mask) without using a SAM prompt encoder + mask decoder.
+            pix_feat = current_vision_feats[-1].permute(1, 2, 0)
+            pix_feat = pix_feat.view(-1, self.hidden_dim, *feat_sizes[-1])
+            sam_outputs = self._use_mask_as_output(
+                pix_feat, high_res_features, mask_inputs
+            )
+        else:
+            # fused the visual feature with previous memory features in the memory bank
+            pix_feat = self._prepare_memory_conditioned_features(
+                frame_idx=frame_idx,
+                is_init_cond_frame=is_init_cond_frame,
+                current_vision_feats=current_vision_feats[-1:],
+                current_vision_pos_embeds=current_vision_pos_embeds[-1:],
+                feat_sizes=feat_sizes[-1:],
+                output_dict=output_dict,
+                num_frames=num_frames,
+                track_in_reverse=track_in_reverse,
+            )
+            # apply SAM-style segmentation head
+            # here we might feed previously predicted low-res SAM mask logits into the SAM mask decoder,
+            # e.g. in demo where such logits come from earlier interaction instead of correction sampling
+            # (in this case, any `mask_inputs` shouldn't reach here as they are sent to _use_mask_as_output instead)
+            if prev_sam_mask_logits is not None:
+                assert point_inputs is not None and mask_inputs is None
+                mask_inputs = prev_sam_mask_logits
+            multimask_output = self._use_multimask(is_init_cond_frame, point_inputs)
+            sam_outputs = self._forward_sam_heads(
+                backbone_features=pix_feat,
+                point_inputs=point_inputs,
+                mask_inputs=mask_inputs,
+                high_res_features=high_res_features,
+                multimask_output=multimask_output,
+            )
+
+        return current_out, sam_outputs, high_res_features, pix_feat
+
+    def _encode_memory_in_output(
+        self,
+        current_vision_feats,
+        feat_sizes,
+        point_inputs,
+        run_mem_encoder,
+        high_res_masks,
+        object_score_logits,
+        current_out,
+    ):
+        if run_mem_encoder and self.num_maskmem > 0:
+            high_res_masks_for_mem_enc = high_res_masks
+            maskmem_features, maskmem_pos_enc = self._encode_new_memory(
+                current_vision_feats=current_vision_feats,
+                feat_sizes=feat_sizes,
+                pred_masks_high_res=high_res_masks_for_mem_enc,
+                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
+
+    def track_step(
+        self,
+        frame_idx,
+        is_init_cond_frame,
+        current_vision_feats,
+        current_vision_pos_embeds,
+        feat_sizes,
+        point_inputs,
+        mask_inputs,
+        output_dict,
+        num_frames,
+        track_in_reverse=False,  # tracking in reverse time order (for demo usage)
+        # Whether to run the memory encoder on the predicted masks. Sometimes we might want
+        # to skip the memory encoder with `run_mem_encoder=False`. For example,
+        # in demo we might call `track_step` multiple times for each user click,
+        # and only encode the memory when the user finalizes their clicks. And in ablation
+        # settings like SAM training on static images, we don't need the memory encoder.
+        run_mem_encoder=True,
+        # The previously predicted SAM mask logits (which can be fed together with new clicks in demo).
+        prev_sam_mask_logits=None,
+    ):
+        current_out, sam_outputs, _, _ = self._track_step(
+            frame_idx,
+            is_init_cond_frame,
+            current_vision_feats,
+            current_vision_pos_embeds,
+            feat_sizes,
+            point_inputs,
+            mask_inputs,
+            output_dict,
+            num_frames,
+            track_in_reverse,
+            prev_sam_mask_logits,
+        )
+
+        (
+            _,
+            _,
+            _,
+            low_res_masks,
+            high_res_masks,
+            obj_ptr,
+            object_score_logits,
+        ) = sam_outputs
+
+        current_out["pred_masks"] = low_res_masks
+        current_out["pred_masks_high_res"] = high_res_masks
+        current_out["obj_ptr"] = obj_ptr
+        if not self.training:
+            # Only add this in inference (to avoid unused param in activation checkpointing;
+            # it's mainly used in the demo to encode spatial memories w/ consolidated masks)
+            current_out["object_score_logits"] = object_score_logits
+
+        # Finally run the memory encoder on the predicted mask to encode
+        # it into a new memory feature (that can be used in future frames)
+        self._encode_memory_in_output(
+            current_vision_feats,
+            feat_sizes,
+            point_inputs,
+            run_mem_encoder,
+            high_res_masks,
+            object_score_logits,
+            current_out,
+        )
+
+        return current_out
+
+    def _use_multimask(self, is_init_cond_frame, point_inputs):
+        """Whether to use multimask output in the SAM head."""
+        num_pts = 0 if point_inputs is None else point_inputs["point_labels"].size(1)
+        multimask_output = (
+            self.multimask_output_in_sam
+            and (is_init_cond_frame or self.multimask_output_for_tracking)
+            and (self.multimask_min_pt_num <= num_pts <= self.multimask_max_pt_num)
+        )
+        return multimask_output
+
+    def _apply_non_overlapping_constraints(self, pred_masks):
+        """
+        Apply non-overlapping constraints to the object scores in pred_masks. Here we
+        keep only the highest scoring object at each spatial location in pred_masks.
+        """
+        batch_size = pred_masks.size(0)
+        if batch_size == 1:
+            return pred_masks
+
+        device = pred_masks.device
+        # "max_obj_inds": object index of the object with the highest score at each location
+        max_obj_inds = torch.argmax(pred_masks, dim=0, keepdim=True)
+        # "batch_obj_inds": object index of each object slice (along dim 0) in `pred_masks`
+        batch_obj_inds = torch.arange(batch_size, device=device)[:, None, None, None]
+        keep = max_obj_inds == batch_obj_inds
+        # suppress overlapping regions' scores below -10.0 so that the foreground regions
+        # don't overlap (here sigmoid(-10.0)=4.5398e-05)
+        pred_masks = torch.where(keep, pred_masks, torch.clamp(pred_masks, max=-10.0))
+        return pred_masks

utility/sam/modeling/sam2_utils.py

@@ -0,0 +1,337 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import copy
+from typing import Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from sam2.utils.misc import mask_to_box
+
+
+def select_closest_cond_frames(frame_idx, cond_frame_outputs, max_cond_frame_num):
+    """
+    Select up to `max_cond_frame_num` conditioning frames from `cond_frame_outputs`
+    that are temporally closest to the current frame at `frame_idx`. Here, we take
+    - a) the closest conditioning frame before `frame_idx` (if any);
+    - b) the closest conditioning frame after `frame_idx` (if any);
+    - c) any other temporally closest conditioning frames until reaching a total
+         of `max_cond_frame_num` conditioning frames.
+
+    Outputs:
+    - selected_outputs: selected items (keys & values) from `cond_frame_outputs`.
+    - unselected_outputs: items (keys & values) not selected in `cond_frame_outputs`.
+    """
+    if max_cond_frame_num == -1 or len(cond_frame_outputs) <= max_cond_frame_num:
+        selected_outputs = cond_frame_outputs
+        unselected_outputs = {}
+    else:
+        assert max_cond_frame_num >= 2, "we should allow using 2+ conditioning frames"
+        selected_outputs = {}
+
+        # the closest conditioning frame before `frame_idx` (if any)
+        idx_before = max((t for t in cond_frame_outputs if t < frame_idx), default=None)
+        if idx_before is not None:
+            selected_outputs[idx_before] = cond_frame_outputs[idx_before]
+
+        # the closest conditioning frame after `frame_idx` (if any)
+        idx_after = min((t for t in cond_frame_outputs if t >= frame_idx), default=None)
+        if idx_after is not None:
+            selected_outputs[idx_after] = cond_frame_outputs[idx_after]
+
+        # add other temporally closest conditioning frames until reaching a total
+        # of `max_cond_frame_num` conditioning frames.
+        num_remain = max_cond_frame_num - len(selected_outputs)
+        inds_remain = sorted(
+            (t for t in cond_frame_outputs if t not in selected_outputs),
+            key=lambda x: abs(x - frame_idx),
+        )[:num_remain]
+        selected_outputs.update((t, cond_frame_outputs[t]) for t in inds_remain)
+        unselected_outputs = {
+            t: v for t, v in cond_frame_outputs.items() if t not in selected_outputs
+        }
+
+    return selected_outputs, unselected_outputs
+
+
+def get_1d_sine_pe(pos_inds, dim, temperature=10000):
+    """
+    Get 1D sine positional embedding as in the original Transformer paper.
+    """
+    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
+    pos_embed = torch.cat([pos_embed.sin(), pos_embed.cos()], dim=-1)
+    return pos_embed
+
+
+def get_activation_fn(activation):
+    """Return an activation function given a string"""
+    if activation == "relu":
+        return F.relu
+    if activation == "gelu":
+        return F.gelu
+    if activation == "glu":
+        return F.glu
+    raise RuntimeError(f"activation should be relu/gelu, not {activation}.")
+
+
+def get_clones(module, N):
+    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])
+
+
+class DropPath(nn.Module):
+    # adapted from https://github.com/huggingface/pytorch-image-models/blob/main/timm/layers/drop.py
+    def __init__(self, drop_prob=0.0, scale_by_keep=True):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+        self.scale_by_keep = scale_by_keep
+
+    def forward(self, x):
+        if self.drop_prob == 0.0 or not self.training:
+            return x
+        keep_prob = 1 - self.drop_prob
+        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+        random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+        if keep_prob > 0.0 and self.scale_by_keep:
+            random_tensor.div_(keep_prob)
+        return x * random_tensor
+
+
+# Lightly adapted from
+# https://github.com/facebookresearch/MaskFormer/blob/main/mask_former/modeling/transformer/transformer_predictor.py # noqa
+class MLP(nn.Module):
+    def __init__(
+        self,
+        input_dim: int,
+        hidden_dim: int,
+        output_dim: int,
+        num_layers: int,
+        activation: nn.Module = nn.ReLU,
+        sigmoid_output: bool = False,
+    ) -> None:
+        super().__init__()
+        self.num_layers = num_layers
+        h = [hidden_dim] * (num_layers - 1)
+        self.layers = nn.ModuleList(
+            nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])
+        )
+        self.sigmoid_output = sigmoid_output
+        self.act = activation()
+
+    def forward(self, x):
+        for i, layer in enumerate(self.layers):
+            x = self.act(layer(x)) if i < self.num_layers - 1 else layer(x)
+        if self.sigmoid_output:
+            x = F.sigmoid(x)
+        return x
+
+
+# From https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/batch_norm.py # noqa
+# Itself from https://github.com/facebookresearch/ConvNeXt/blob/d1fa8f6fef0a165b27399986cc2bdacc92777e40/models/convnext.py#L119  # noqa
+class LayerNorm2d(nn.Module):
+    def __init__(self, num_channels: int, eps: float = 1e-6) -> None:
+        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: torch.Tensor) -> torch.Tensor:
+        u = x.mean(1, keepdim=True)
+        s = (x - u).pow(2).mean(1, keepdim=True)
+        x = (x - u) / torch.sqrt(s + self.eps)
+        x = self.weight[:, None, None] * x + self.bias[:, None, None]
+        return x
+
+
+def sample_box_points(
+    masks: torch.Tensor,
+    noise: float = 0.1,  # SAM default
+    noise_bound: int = 20,  # SAM default
+    top_left_label: int = 2,
+    bottom_right_label: int = 3,
+) -> Tuple[np.array, np.array]:
+    """
+    Sample a noised version of the top left and bottom right corners of a given `bbox`
+
+    Inputs:
+    - masks: [B, 1, H,W] boxes, dtype=torch.Tensor
+    - noise: noise as a fraction of box width and height, dtype=float
+    - noise_bound: maximum amount of noise (in pure pixesl), dtype=int
+
+    Returns:
+    - box_coords: [B, num_pt, 2], contains (x, y) coordinates of top left and bottom right box corners, dtype=torch.float
+    - box_labels: [B, num_pt], label 2 is reserverd for top left and 3 for bottom right corners, dtype=torch.int32
+    """
+    device = masks.device
+    box_coords = mask_to_box(masks)
+    B, _, H, W = masks.shape
+    box_labels = torch.tensor(
+        [top_left_label, bottom_right_label], dtype=torch.int, device=device
+    ).repeat(B)
+    if noise > 0.0:
+        if not isinstance(noise_bound, torch.Tensor):
+            noise_bound = torch.tensor(noise_bound, device=device)
+        bbox_w = box_coords[..., 2] - box_coords[..., 0]
+        bbox_h = box_coords[..., 3] - box_coords[..., 1]
+        max_dx = torch.min(bbox_w * noise, noise_bound)
+        max_dy = torch.min(bbox_h * noise, noise_bound)
+        box_noise = 2 * torch.rand(B, 1, 4, device=device) - 1
+        box_noise = box_noise * torch.stack((max_dx, max_dy, max_dx, max_dy), dim=-1)
+
+        box_coords = box_coords + box_noise
+        img_bounds = (
+            torch.tensor([W, H, W, H], device=device) - 1
+        )  # uncentered pixel coords
+        box_coords.clamp_(torch.zeros_like(img_bounds), img_bounds)  # In place clamping
+
+    box_coords = box_coords.reshape(-1, 2, 2)  # always 2 points
+    box_labels = box_labels.reshape(-1, 2)
+    return box_coords, box_labels
+
+
+def sample_random_points_from_errors(gt_masks, pred_masks, num_pt=1):
+    """
+    Sample `num_pt` random points (along with their labels) independently from the error regions.
+
+    Inputs:
+    - gt_masks: [B, 1, H_im, W_im] masks, dtype=torch.bool
+    - pred_masks: [B, 1, H_im, W_im] masks, dtype=torch.bool or None
+    - num_pt: int, number of points to sample independently for each of the B error maps
+
+    Outputs:
+    - points: [B, num_pt, 2], dtype=torch.float, contains (x, y) coordinates of each sampled point
+    - labels: [B, num_pt], dtype=torch.int32, where 1 means positive clicks and 0 means
+      negative clicks
+    """
+    if pred_masks is None:  # if pred_masks is not provided, treat it as empty
+        pred_masks = torch.zeros_like(gt_masks)
+    assert gt_masks.dtype == torch.bool and gt_masks.size(1) == 1
+    assert pred_masks.dtype == torch.bool and pred_masks.shape == gt_masks.shape
+    assert num_pt >= 0
+
+    B, _, H_im, W_im = gt_masks.shape
+    device = gt_masks.device
+
+    # false positive region, a new point sampled in this region should have
+    # negative label to correct the FP error
+    fp_masks = ~gt_masks & pred_masks
+    # false negative region, a new point sampled in this region should have
+    # positive label to correct the FN error
+    fn_masks = gt_masks & ~pred_masks
+    # whether the prediction completely match the ground-truth on each mask
+    all_correct = torch.all((gt_masks == pred_masks).flatten(2), dim=2)
+    all_correct = all_correct[..., None, None]
+
+    # channel 0 is FP map, while channel 1 is FN map
+    pts_noise = torch.rand(B, num_pt, H_im, W_im, 2, device=device)
+    # sample a negative new click from FP region or a positive new click
+    # from FN region, depend on where the maximum falls,
+    # and in case the predictions are all correct (no FP or FN), we just
+    # sample a negative click from the background region
+    pts_noise[..., 0] *= fp_masks | (all_correct & ~gt_masks)
+    pts_noise[..., 1] *= fn_masks
+    pts_idx = pts_noise.flatten(2).argmax(dim=2)
+    labels = (pts_idx % 2).to(torch.int32)
+    pts_idx = pts_idx // 2
+    pts_x = pts_idx % W_im
+    pts_y = pts_idx // W_im
+    points = torch.stack([pts_x, pts_y], dim=2).to(torch.float)
+    return points, labels
+
+
+def sample_one_point_from_error_center(gt_masks, pred_masks, padding=True):
+    """
+    Sample 1 random point (along with its label) from the center of each error region,
+    that is, the point with the largest distance to the boundary of each error region.
+    This is the RITM sampling method from https://github.com/saic-vul/ritm_interactive_segmentation/blob/master/isegm/inference/clicker.py
+
+    Inputs:
+    - gt_masks: [B, 1, H_im, W_im] masks, dtype=torch.bool
+    - pred_masks: [B, 1, H_im, W_im] masks, dtype=torch.bool or None
+    - padding: if True, pad with boundary of 1 px for distance transform
+
+    Outputs:
+    - points: [B, 1, 2], dtype=torch.float, contains (x, y) coordinates of each sampled point
+    - labels: [B, 1], dtype=torch.int32, where 1 means positive clicks and 0 means negative clicks
+    """
+    import cv2
+
+    if pred_masks is None:
+        pred_masks = torch.zeros_like(gt_masks)
+    assert gt_masks.dtype == torch.bool and gt_masks.size(1) == 1
+    assert pred_masks.dtype == torch.bool and pred_masks.shape == gt_masks.shape
+
+    B, _, _, W_im = gt_masks.shape
+    device = gt_masks.device
+
+    # false positive region, a new point sampled in this region should have
+    # negative label to correct the FP error
+    fp_masks = ~gt_masks & pred_masks
+    # false negative region, a new point sampled in this region should have
+    # positive label to correct the FN error
+    fn_masks = gt_masks & ~pred_masks
+
+    fp_masks = fp_masks.cpu().numpy()
+    fn_masks = fn_masks.cpu().numpy()
+    points = torch.zeros(B, 1, 2, dtype=torch.float)
+    labels = torch.ones(B, 1, dtype=torch.int32)
+    for b in range(B):
+        fn_mask = fn_masks[b, 0]
+        fp_mask = fp_masks[b, 0]
+        if padding:
+            fn_mask = np.pad(fn_mask, ((1, 1), (1, 1)), "constant")
+            fp_mask = np.pad(fp_mask, ((1, 1), (1, 1)), "constant")
+        # compute the distance of each point in FN/FP region to its boundary
+        fn_mask_dt = cv2.distanceTransform(fn_mask.astype(np.uint8), cv2.DIST_L2, 0)
+        fp_mask_dt = cv2.distanceTransform(fp_mask.astype(np.uint8), cv2.DIST_L2, 0)
+        if padding:
+            fn_mask_dt = fn_mask_dt[1:-1, 1:-1]
+            fp_mask_dt = fp_mask_dt[1:-1, 1:-1]
+
+        # take the point in FN/FP region with the largest distance to its boundary
+        fn_mask_dt_flat = fn_mask_dt.reshape(-1)
+        fp_mask_dt_flat = fp_mask_dt.reshape(-1)
+        fn_argmax = np.argmax(fn_mask_dt_flat)
+        fp_argmax = np.argmax(fp_mask_dt_flat)
+        is_positive = fn_mask_dt_flat[fn_argmax] > fp_mask_dt_flat[fp_argmax]
+        pt_idx = fn_argmax if is_positive else fp_argmax
+        points[b, 0, 0] = pt_idx % W_im  # x
+        points[b, 0, 1] = pt_idx // W_im  # y
+        labels[b, 0] = int(is_positive)
+
+    points = points.to(device)
+    labels = labels.to(device)
+    return points, labels
+
+
+def get_next_point(gt_masks, pred_masks, method):
+    if method == "uniform":
+        return sample_random_points_from_errors(gt_masks, pred_masks)
+    elif method == "center":
+        return sample_one_point_from_error_center(gt_masks, pred_masks)
+    else:
+        raise ValueError(f"unknown sampling method {method}")
+
+class LayerScale(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        init_values: Union[float, torch.Tensor] = 1e-5,
+        inplace: bool = False,
+    ) -> None:
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma

utility/sam/sam2_image_predictor.py

@@ -0,0 +1,461 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from PIL.Image import Image
+
+from sam2.modeling.sam2_base import SAM2Base
+
+from sam2.utils.transforms import SAM2Transforms
+
+
+class SAM2ImagePredictor:
+    def __init__(
+        self,
+        sam_model: SAM2Base,
+        mask_threshold=0.0,
+        max_hole_area=0.0,
+        max_sprinkle_area=0.0,
+        **kwargs,
+    ) -> None:
+        """
+        Uses SAM-2 to calculate the image embedding for an image, and then
+        allow repeated, efficient mask prediction given prompts.
+
+        Arguments:
+          sam_model (Sam-2): The model to use for mask prediction.
+          mask_threshold (float): The threshold to use when converting mask logits
+            to binary masks. Masks are thresholded at 0 by default.
+          max_hole_area (int): If max_hole_area > 0, we fill small holes in up to
+            the maximum area of max_hole_area in low_res_masks.
+          max_sprinkle_area (int): If max_sprinkle_area > 0, we remove small sprinkles up to
+            the maximum area of max_sprinkle_area in low_res_masks.
+        """
+        super().__init__()
+        self.model = sam_model
+        self._transforms = SAM2Transforms(
+            resolution=self.model.image_size,
+            mask_threshold=mask_threshold,
+            max_hole_area=max_hole_area,
+            max_sprinkle_area=max_sprinkle_area,
+        )
+
+        # Predictor state
+        self._is_image_set = False
+        self._features = None
+        self._orig_hw = None
+        # Whether the predictor is set for single image or a batch of images
+        self._is_batch = False
+
+        # Predictor config
+        self.mask_threshold = mask_threshold
+
+        # Spatial dim for backbone feature maps
+        self._bb_feat_sizes = [
+            (256, 256),
+            (128, 128),
+            (64, 64),
+        ]
+
+    @classmethod
+    def from_pretrained(cls, model_id: str, **kwargs) -> "SAM2ImagePredictor":
+        """
+        Load a pretrained model from the Hugging Face hub.
+
+        Arguments:
+          model_id (str): The Hugging Face repository ID.
+          **kwargs: Additional arguments to pass to the model constructor.
+
+        Returns:
+          (SAM2ImagePredictor): The loaded model.
+        """
+        from sam2.build_sam import build_sam2_hf
+
+        sam_model = build_sam2_hf(model_id, **kwargs)
+        return cls(sam_model, **kwargs)
+
+    @torch.no_grad()
+    def set_image(
+        self,
+        image: Union[np.ndarray, Image],
+    ) -> None:
+        """
+        Calculates the image embeddings for the provided image, allowing
+        masks to be predicted with the 'predict' method.
+
+        Arguments:
+          image (np.ndarray or PIL Image): The input image to embed in RGB format. The image should be in HWC format if np.ndarray, or WHC format if PIL Image
+          with pixel values in [0, 255].
+          image_format (str): The color format of the image, in ['RGB', 'BGR'].
+        """
+        self.reset_predictor()
+        # Transform the image to the form expected by the model
+        if isinstance(image, np.ndarray):
+            self._orig_hw = [image.shape[:2]]
+        elif isinstance(image, Image):
+            w, h = image.size
+            self._orig_hw = [(h, w)]
+        else:
+            raise NotImplementedError("Image format not supported")
+
+        input_image = self._transforms(image)
+        input_image = input_image[None, ...].to(self.device)
+
+        assert (
+            len(input_image.shape) == 4 and input_image.shape[1] == 3
+        ), f"input_image must be of size 1x3xHxW, got {input_image.shape}"
+        backbone_out = self.model.forward_image(input_image)
+        _, vision_feats, _, _ = self.model._prepare_backbone_features(backbone_out)
+        # Add no_mem_embed, which is added to the lowest rest feat. map during training on videos
+        if self.model.directly_add_no_mem_embed:
+            vision_feats[-1] = vision_feats[-1] + self.model.no_mem_embed
+
+        feats = [
+            feat.permute(1, 2, 0).view(1, -1, *feat_size)
+            for feat, feat_size in zip(vision_feats[::-1], self._bb_feat_sizes[::-1])
+        ][::-1]
+        self._features = {"image_embed": feats[-1], "high_res_feats": feats[:-1]}
+        self._is_image_set = True
+
+    @torch.no_grad()
+    def set_image_batch(
+        self,
+        image_list: List[Union[np.ndarray]],
+    ) -> None:
+        """
+        Calculates the image embeddings for the provided image batch, allowing
+        masks to be predicted with the 'predict_batch' method.
+
+        Arguments:
+          image_list (List[np.ndarray]): The input images to embed in RGB format. The image should be in HWC format if np.ndarray
+          with pixel values in [0, 255].
+        """
+        self.reset_predictor()
+        assert isinstance(image_list, list)
+        self._orig_hw = []
+        for image in image_list:
+            assert isinstance(
+                image, np.ndarray
+            ), "Images are expected to be an np.ndarray in RGB format, and of shape  HWC"
+            self._orig_hw.append(image.shape[:2])
+        # Transform the image to the form expected by the model
+        img_batch = self._transforms.forward_batch(image_list)
+        img_batch = img_batch.to(self.device)
+        batch_size = img_batch.shape[0]
+        assert (
+            len(img_batch.shape) == 4 and img_batch.shape[1] == 3
+        ), f"img_batch must be of size Bx3xHxW, got {img_batch.shape}"
+        backbone_out = self.model.forward_image(img_batch)
+        _, vision_feats, _, _ = self.model._prepare_backbone_features(backbone_out)
+        # Add no_mem_embed, which is added to the lowest rest feat. map during training on videos
+        if self.model.directly_add_no_mem_embed:
+            vision_feats[-1] = vision_feats[-1] + self.model.no_mem_embed
+
+        feats = [
+            feat.permute(1, 2, 0).view(batch_size, -1, *feat_size)
+            for feat, feat_size in zip(vision_feats[::-1], self._bb_feat_sizes[::-1])
+        ][::-1]
+        self._features = {"image_embed": feats[-1], "high_res_feats": feats[:-1]}
+        self._is_image_set = True
+        self._is_batch = True
+
+    def predict_batch(
+        self,
+        point_coords_batch: List[np.ndarray] = None,
+        point_labels_batch: List[np.ndarray] = None,
+        box_batch: List[np.ndarray] = None,
+        mask_input_batch: List[np.ndarray] = None,
+        multimask_output: bool = True,
+        return_logits: bool = False,
+        normalize_coords=True,
+    ) -> Tuple[List[np.ndarray], List[np.ndarray], List[np.ndarray]]:
+        """This function is very similar to predict(...), however it is used for batched mode, when the model is expected to generate predictions on multiple images.
+        It returns a tuple of lists of masks, ious, and low_res_masks_logits.
+        """
+        assert self._is_batch, "This function should only be used when in batched mode"
+        if not self._is_image_set:
+            raise RuntimeError(
+                "An image must be set with .set_image_batch(...) before mask prediction."
+            )
+        num_images = len(self._features["image_embed"])
+        all_masks = []
+        all_ious = []
+        all_low_res_masks = []
+        for img_idx in range(num_images):
+            # Transform input prompts
+            point_coords = (
+                point_coords_batch[img_idx] if point_coords_batch is not None else None
+            )
+            point_labels = (
+                point_labels_batch[img_idx] if point_labels_batch is not None else None
+            )
+            box = box_batch[img_idx] if box_batch is not None else None
+            mask_input = (
+                mask_input_batch[img_idx] if mask_input_batch is not None else None
+            )
+            mask_input, unnorm_coords, labels, unnorm_box = self._prep_prompts(
+                point_coords,
+                point_labels,
+                box,
+                mask_input,
+                normalize_coords,
+                img_idx=img_idx,
+            )
+            masks, iou_predictions, low_res_masks = self._predict(
+                unnorm_coords,
+                labels,
+                unnorm_box,
+                mask_input,
+                multimask_output,
+                return_logits=return_logits,
+                img_idx=img_idx,
+            )
+            masks_np = masks.squeeze(0).float().detach().cpu().numpy()
+            iou_predictions_np = (
+                iou_predictions.squeeze(0).float().detach().cpu().numpy()
+            )
+            low_res_masks_np = low_res_masks.squeeze(0).float().detach().cpu().numpy()
+            all_masks.append(masks_np)
+            all_ious.append(iou_predictions_np)
+            all_low_res_masks.append(low_res_masks_np)
+
+        return all_masks, all_ious, all_low_res_masks
+
+    def predict(
+        self,
+        point_coords: Optional[np.ndarray] = None,
+        point_labels: Optional[np.ndarray] = None,
+        box: Optional[np.ndarray] = None,
+        mask_input: Optional[np.ndarray] = None,
+        multimask_output: bool = True,
+        return_logits: bool = False,
+        normalize_coords=True,
+    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """
+        Predict masks for the given input prompts, using the currently set image.
+
+        Arguments:
+          point_coords (np.ndarray or None): A Nx2 array of point prompts to the
+            model. Each point is in (X,Y) in pixels.
+          point_labels (np.ndarray or None): A length N array of labels for the
+            point prompts. 1 indicates a foreground point and 0 indicates a
+            background point.
+          box (np.ndarray or None): A length 4 array given a box prompt to the
+            model, in XYXY format.
+          mask_input (np.ndarray): A low resolution mask input to the model, typically
+            coming from a previous prediction iteration. Has form 1xHxW, where
+            for SAM, H=W=256.
+          multimask_output (bool): If true, the model will return three masks.
+            For ambiguous input prompts (such as a single click), this will often
+            produce better masks than a single prediction. If only a single
+            mask is needed, the model's predicted quality score can be used
+            to select the best mask. For non-ambiguous prompts, such as multiple
+            input prompts, multimask_output=False can give better results.
+          return_logits (bool): If true, returns un-thresholded masks logits
+            instead of a binary mask.
+          normalize_coords (bool): If true, the point coordinates will be normalized to the range [0,1] and point_coords is expected to be wrt. image dimensions.
+
+        Returns:
+          (np.ndarray): The output masks in CxHxW format, where C is the
+            number of masks, and (H, W) is the original image size.
+          (np.ndarray): An array of length C containing the model's
+            predictions for the quality of each mask.
+          (np.ndarray): An array of shape CxHxW, where C is the number
+            of masks and H=W=256. These low resolution logits can be passed to
+            a subsequent iteration as mask input.
+        """
+        if not self._is_image_set:
+            raise RuntimeError(
+                "An image must be set with .set_image(...) before mask prediction."
+            )
+
+        # Transform input prompts
+
+        mask_input, unnorm_coords, labels, unnorm_box = self._prep_prompts(
+            point_coords, point_labels, box, mask_input, normalize_coords
+        )
+
+        masks, iou_predictions, low_res_masks = self._predict(
+            unnorm_coords,
+            labels,
+            unnorm_box,
+            mask_input,
+            multimask_output,
+            return_logits=return_logits,
+        )
+
+        masks_np = masks.squeeze(0).float().detach().cpu().numpy()
+        iou_predictions_np = iou_predictions.squeeze(0).float().detach().cpu().numpy()
+        low_res_masks_np = low_res_masks.squeeze(0).float().detach().cpu().numpy()
+        return masks_np, iou_predictions_np, low_res_masks_np
+
+    def _prep_prompts(
+        self, point_coords, point_labels, box, mask_logits, normalize_coords, img_idx=-1
+    ):
+
+        unnorm_coords, labels, unnorm_box, mask_input = None, None, None, None
+        if point_coords is not None:
+            assert (
+                point_labels is not None
+            ), "point_labels must be supplied if point_coords is supplied."
+            point_coords = torch.as_tensor(
+                point_coords, dtype=torch.float, device=self.device
+            )
+            unnorm_coords = self._transforms.transform_coords(
+                point_coords, normalize=normalize_coords, orig_hw=self._orig_hw[img_idx]
+            )
+            labels = torch.as_tensor(point_labels, dtype=torch.int, device=self.device)
+            if len(unnorm_coords.shape) == 2:
+                unnorm_coords, labels = unnorm_coords[None, ...], labels[None, ...]
+        if box is not None:
+            box = torch.as_tensor(box, dtype=torch.float, device=self.device)
+            unnorm_box = self._transforms.transform_boxes(
+                box, normalize=normalize_coords, orig_hw=self._orig_hw[img_idx]
+            )  # Bx2x2
+        if mask_logits is not None:
+            mask_input = torch.as_tensor(
+                mask_logits, dtype=torch.float, device=self.device
+            )
+            if len(mask_input.shape) == 3:
+                mask_input = mask_input[None, :, :, :]
+        return mask_input, unnorm_coords, labels, unnorm_box
+
+    @torch.no_grad()
+    def _predict(
+        self,
+        point_coords: Optional[torch.Tensor],
+        point_labels: Optional[torch.Tensor],
+        boxes: Optional[torch.Tensor] = None,
+        mask_input: Optional[torch.Tensor] = None,
+        multimask_output: bool = True,
+        return_logits: bool = False,
+        img_idx: int = -1,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Predict masks for the given input prompts, using the currently set image.
+        Input prompts are batched torch tensors and are expected to already be
+        transformed to the input frame using SAM2Transforms.
+
+        Arguments:
+          point_coords (torch.Tensor or None): A BxNx2 array of point prompts to the
+            model. Each point is in (X,Y) in pixels.
+          point_labels (torch.Tensor or None): A BxN array of labels for the
+            point prompts. 1 indicates a foreground point and 0 indicates a
+            background point.
+          boxes (np.ndarray or None): A Bx4 array given a box prompt to the
+            model, in XYXY format.
+          mask_input (np.ndarray): A low resolution mask input to the model, typically
+            coming from a previous prediction iteration. Has form Bx1xHxW, where
+            for SAM, H=W=256. Masks returned by a previous iteration of the
+            predict method do not need further transformation.
+          multimask_output (bool): If true, the model will return three masks.
+            For ambiguous input prompts (such as a single click), this will often
+            produce better masks than a single prediction. If only a single
+            mask is needed, the model's predicted quality score can be used
+            to select the best mask. For non-ambiguous prompts, such as multiple
+            input prompts, multimask_output=False can give better results.
+          return_logits (bool): If true, returns un-thresholded masks logits
+            instead of a binary mask.
+
+        Returns:
+          (torch.Tensor): The output masks in BxCxHxW format, where C is the
+            number of masks, and (H, W) is the original image size.
+          (torch.Tensor): An array of shape BxC containing the model's
+            predictions for the quality of each mask.
+          (torch.Tensor): An array of shape BxCxHxW, where C is the number
+            of masks and H=W=256. These low res logits can be passed to
+            a subsequent iteration as mask input.
+        """
+        if not self._is_image_set:
+            raise RuntimeError(
+                "An image must be set with .set_image(...) before mask prediction."
+            )
+
+        if point_coords is not None:
+            concat_points = (point_coords, point_labels)
+        else:
+            concat_points = None
+
+        # Embed prompts
+        if boxes is not None:
+            box_coords = boxes.reshape(-1, 2, 2)
+            box_labels = torch.tensor([[2, 3]], dtype=torch.int, device=boxes.device)
+            box_labels = box_labels.repeat(boxes.size(0), 1)
+            # we merge "boxes" and "points" into a single "concat_points" input (where
+            # boxes are added at the beginning) to sam_prompt_encoder
+            if concat_points is not None:
+                concat_coords = torch.cat([box_coords, concat_points[0]], dim=1)
+                concat_labels = torch.cat([box_labels, concat_points[1]], dim=1)
+                concat_points = (concat_coords, concat_labels)
+            else:
+                concat_points = (box_coords, box_labels)
+
+        sparse_embeddings, dense_embeddings = self.model.sam_prompt_encoder(
+            points=concat_points,
+            boxes=None,
+            masks=mask_input,
+        )
+
+        # Predict masks
+        batched_mode = (
+            concat_points is not None and concat_points[0].shape[0] > 1
+        )  # multi object prediction
+        high_res_features = [
+            feat_level[img_idx].unsqueeze(0)
+            for feat_level in self._features["high_res_feats"]
+        ]
+        low_res_masks, iou_predictions, _, _ = self.model.sam_mask_decoder(
+            image_embeddings=self._features["image_embed"][img_idx].unsqueeze(0),
+            image_pe=self.model.sam_prompt_encoder.get_dense_pe(),
+            sparse_prompt_embeddings=sparse_embeddings,
+            dense_prompt_embeddings=dense_embeddings,
+            multimask_output=multimask_output,
+            repeat_image=batched_mode,
+            high_res_features=high_res_features,
+        )
+
+        # Upscale the masks to the original image resolution
+        masks = self._transforms.postprocess_masks(
+            low_res_masks, self._orig_hw[img_idx]
+        )
+        low_res_masks = torch.clamp(low_res_masks, -32.0, 32.0)
+        if not return_logits:
+            masks = masks > self.mask_threshold
+
+        return masks, iou_predictions, low_res_masks
+
+    def get_image_embedding(self) -> torch.Tensor:
+        """
+        Returns the image embeddings for the currently set image, with
+        shape 1xCxHxW, where C is the embedding dimension and (H,W) are
+        the embedding spatial dimension of SAM (typically C=256, H=W=64).
+        """
+        if not self._is_image_set:
+            raise RuntimeError(
+                "An image must be set with .set_image(...) to generate an embedding."
+            )
+        assert (
+            self._features is not None
+        ), "Features must exist if an image has been set."
+        return self._features["image_embed"]
+
+    @property
+    def device(self) -> torch.device:
+        return self.model.device
+
+    def reset_predictor(self) -> None:
+        """
+        Resets the image embeddings and other state variables.
+        """
+        self._is_image_set = False
+        self._features = None
+        self._orig_hw = None
+        self._is_batch = False

utility/sam/utils/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.

utility/sam/utils/amg.py

@@ -0,0 +1,348 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+from copy import deepcopy
+from itertools import product
+from typing import Any, Dict, Generator, ItemsView, List, Tuple
+
+import numpy as np
+import torch
+
+# Very lightly adapted from https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/utils/amg.py
+
+
+class MaskData:
+    """
+    A structure for storing masks and their related data in batched format.
+    Implements basic filtering and concatenation.
+    """
+
+    def __init__(self, **kwargs) -> None:
+        for v in kwargs.values():
+            assert isinstance(
+                v, (list, np.ndarray, torch.Tensor)
+            ), "MaskData only supports list, numpy arrays, and torch tensors."
+        self._stats = dict(**kwargs)
+
+    def __setitem__(self, key: str, item: Any) -> None:
+        assert isinstance(
+            item, (list, np.ndarray, torch.Tensor)
+        ), "MaskData only supports list, numpy arrays, and torch tensors."
+        self._stats[key] = item
+
+    def __delitem__(self, key: str) -> None:
+        del self._stats[key]
+
+    def __getitem__(self, key: str) -> Any:
+        return self._stats[key]
+
+    def items(self) -> ItemsView[str, Any]:
+        return self._stats.items()
+
+    def filter(self, keep: torch.Tensor) -> None:
+        for k, v in self._stats.items():
+            if v is None:
+                self._stats[k] = None
+            elif isinstance(v, torch.Tensor):
+                self._stats[k] = v[torch.as_tensor(keep, device=v.device)]
+            elif isinstance(v, np.ndarray):
+                self._stats[k] = v[keep.detach().cpu().numpy()]
+            elif isinstance(v, list) and keep.dtype == torch.bool:
+                self._stats[k] = [a for i, a in enumerate(v) if keep[i]]
+            elif isinstance(v, list):
+                self._stats[k] = [v[i] for i in keep]
+            else:
+                raise TypeError(f"MaskData key {k} has an unsupported type {type(v)}.")
+
+    def cat(self, new_stats: "MaskData") -> None:
+        for k, v in new_stats.items():
+            if k not in self._stats or self._stats[k] is None:
+                self._stats[k] = deepcopy(v)
+            elif isinstance(v, torch.Tensor):
+                self._stats[k] = torch.cat([self._stats[k], v], dim=0)
+            elif isinstance(v, np.ndarray):
+                self._stats[k] = np.concatenate([self._stats[k], v], axis=0)
+            elif isinstance(v, list):
+                self._stats[k] = self._stats[k] + deepcopy(v)
+            else:
+                raise TypeError(f"MaskData key {k} has an unsupported type {type(v)}.")
+
+    def to_numpy(self) -> None:
+        for k, v in self._stats.items():
+            if isinstance(v, torch.Tensor):
+                self._stats[k] = v.float().detach().cpu().numpy()
+
+
+def is_box_near_crop_edge(
+    boxes: torch.Tensor, crop_box: List[int], orig_box: List[int], atol: float = 20.0
+) -> torch.Tensor:
+    """Filter masks at the edge of a crop, but not at the edge of the original image."""
+    crop_box_torch = torch.as_tensor(crop_box, dtype=torch.float, device=boxes.device)
+    orig_box_torch = torch.as_tensor(orig_box, dtype=torch.float, device=boxes.device)
+    boxes = uncrop_boxes_xyxy(boxes, crop_box).float()
+    near_crop_edge = torch.isclose(boxes, crop_box_torch[None, :], atol=atol, rtol=0)
+    near_image_edge = torch.isclose(boxes, orig_box_torch[None, :], atol=atol, rtol=0)
+    near_crop_edge = torch.logical_and(near_crop_edge, ~near_image_edge)
+    return torch.any(near_crop_edge, dim=1)
+
+
+def box_xyxy_to_xywh(box_xyxy: torch.Tensor) -> torch.Tensor:
+    box_xywh = deepcopy(box_xyxy)
+    box_xywh[2] = box_xywh[2] - box_xywh[0]
+    box_xywh[3] = box_xywh[3] - box_xywh[1]
+    return box_xywh
+
+
+def batch_iterator(batch_size: int, *args) -> Generator[List[Any], None, None]:
+    assert len(args) > 0 and all(
+        len(a) == len(args[0]) for a in args
+    ), "Batched iteration must have inputs of all the same size."
+    n_batches = len(args[0]) // batch_size + int(len(args[0]) % batch_size != 0)
+    for b in range(n_batches):
+        yield [arg[b * batch_size : (b + 1) * batch_size] for arg in args]
+
+
+def mask_to_rle_pytorch(tensor: torch.Tensor) -> List[Dict[str, Any]]:
+    """
+    Encodes masks to an uncompressed RLE, in the format expected by
+    pycoco tools.
+    """
+    # Put in fortran order and flatten h,w
+    b, h, w = tensor.shape
+    tensor = tensor.permute(0, 2, 1).flatten(1)
+
+    # Compute change indices
+    diff = tensor[:, 1:] ^ tensor[:, :-1]
+    change_indices = diff.nonzero()
+
+    # Encode run length
+    out = []
+    for i in range(b):
+        cur_idxs = change_indices[change_indices[:, 0] == i, 1]
+        cur_idxs = torch.cat(
+            [
+                torch.tensor([0], dtype=cur_idxs.dtype, device=cur_idxs.device),
+                cur_idxs + 1,
+                torch.tensor([h * w], dtype=cur_idxs.dtype, device=cur_idxs.device),
+            ]
+        )
+        btw_idxs = cur_idxs[1:] - cur_idxs[:-1]
+        counts = [] if tensor[i, 0] == 0 else [0]
+        counts.extend(btw_idxs.detach().cpu().tolist())
+        out.append({"size": [h, w], "counts": counts})
+    return out
+
+
+def rle_to_mask(rle: Dict[str, Any]) -> np.ndarray:
+    """Compute a binary mask from an uncompressed RLE."""
+    h, w = rle["size"]
+    mask = np.empty(h * w, dtype=bool)
+    idx = 0
+    parity = False
+    for count in rle["counts"]:
+        mask[idx : idx + count] = parity
+        idx += count
+        parity ^= True
+    mask = mask.reshape(w, h)
+    return mask.transpose()  # Put in C order
+
+
+def area_from_rle(rle: Dict[str, Any]) -> int:
+    return sum(rle["counts"][1::2])
+
+
+def calculate_stability_score(
+    masks: torch.Tensor, mask_threshold: float, threshold_offset: float
+) -> torch.Tensor:
+    """
+    Computes the stability score for a batch of masks. The stability
+    score is the IoU between the binary masks obtained by thresholding
+    the predicted mask logits at high and low values.
+    """
+    # One mask is always contained inside the other.
+    # Save memory by preventing unnecessary cast to torch.int64
+    intersections = (
+        (masks > (mask_threshold + threshold_offset))
+        .sum(-1, dtype=torch.int16)
+        .sum(-1, dtype=torch.int32)
+    )
+    unions = (
+        (masks > (mask_threshold - threshold_offset))
+        .sum(-1, dtype=torch.int16)
+        .sum(-1, dtype=torch.int32)
+    )
+    return intersections / unions
+
+
+def build_point_grid(n_per_side: int) -> np.ndarray:
+    """Generates a 2D grid of points evenly spaced in [0,1]x[0,1]."""
+    offset = 1 / (2 * n_per_side)
+    points_one_side = np.linspace(offset, 1 - offset, n_per_side)
+    points_x = np.tile(points_one_side[None, :], (n_per_side, 1))
+    points_y = np.tile(points_one_side[:, None], (1, n_per_side))
+    points = np.stack([points_x, points_y], axis=-1).reshape(-1, 2)
+    return points
+
+
+def build_all_layer_point_grids(
+    n_per_side: int, n_layers: int, scale_per_layer: int
+) -> List[np.ndarray]:
+    """Generates point grids for all crop layers."""
+    points_by_layer = []
+    for i in range(n_layers + 1):
+        n_points = int(n_per_side / (scale_per_layer**i))
+        points_by_layer.append(build_point_grid(n_points))
+    return points_by_layer
+
+
+def generate_crop_boxes(
+    im_size: Tuple[int, ...], n_layers: int, overlap_ratio: float
+) -> Tuple[List[List[int]], List[int]]:
+    """
+    Generates a list of crop boxes of different sizes. Each layer
+    has (2**i)**2 boxes for the ith layer.
+    """
+    crop_boxes, layer_idxs = [], []
+    im_h, im_w = im_size
+    short_side = min(im_h, im_w)
+
+    # Original image
+    crop_boxes.append([0, 0, im_w, im_h])
+    layer_idxs.append(0)
+
+    def crop_len(orig_len, n_crops, overlap):
+        return int(math.ceil((overlap * (n_crops - 1) + orig_len) / n_crops))
+
+    for i_layer in range(n_layers):
+        n_crops_per_side = 2 ** (i_layer + 1)
+        overlap = int(overlap_ratio * short_side * (2 / n_crops_per_side))
+
+        crop_w = crop_len(im_w, n_crops_per_side, overlap)
+        crop_h = crop_len(im_h, n_crops_per_side, overlap)
+
+        crop_box_x0 = [int((crop_w - overlap) * i) for i in range(n_crops_per_side)]
+        crop_box_y0 = [int((crop_h - overlap) * i) for i in range(n_crops_per_side)]
+
+        # Crops in XYWH format
+        for x0, y0 in product(crop_box_x0, crop_box_y0):
+            box = [x0, y0, min(x0 + crop_w, im_w), min(y0 + crop_h, im_h)]
+            crop_boxes.append(box)
+            layer_idxs.append(i_layer + 1)
+
+    return crop_boxes, layer_idxs
+
+
+def uncrop_boxes_xyxy(boxes: torch.Tensor, crop_box: List[int]) -> torch.Tensor:
+    x0, y0, _, _ = crop_box
+    offset = torch.tensor([[x0, y0, x0, y0]], device=boxes.device)
+    # Check if boxes has a channel dimension
+    if len(boxes.shape) == 3:
+        offset = offset.unsqueeze(1)
+    return boxes + offset
+
+
+def uncrop_points(points: torch.Tensor, crop_box: List[int]) -> torch.Tensor:
+    x0, y0, _, _ = crop_box
+    offset = torch.tensor([[x0, y0]], device=points.device)
+    # Check if points has a channel dimension
+    if len(points.shape) == 3:
+        offset = offset.unsqueeze(1)
+    return points + offset
+
+
+def uncrop_masks(
+    masks: torch.Tensor, crop_box: List[int], orig_h: int, orig_w: int
+) -> torch.Tensor:
+    x0, y0, x1, y1 = crop_box
+    if x0 == 0 and y0 == 0 and x1 == orig_w and y1 == orig_h:
+        return masks
+    # Coordinate transform masks
+    pad_x, pad_y = orig_w - (x1 - x0), orig_h - (y1 - y0)
+    pad = (x0, pad_x - x0, y0, pad_y - y0)
+    return torch.nn.functional.pad(masks, pad, value=0)
+
+
+def remove_small_regions(
+    mask: np.ndarray, area_thresh: float, mode: str
+) -> Tuple[np.ndarray, bool]:
+    """
+    Removes small disconnected regions and holes in a mask. Returns the
+    mask and an indicator of if the mask has been modified.
+    """
+    import cv2  # type: ignore
+
+    assert mode in ["holes", "islands"]
+    correct_holes = mode == "holes"
+    working_mask = (correct_holes ^ mask).astype(np.uint8)
+    n_labels, regions, stats, _ = cv2.connectedComponentsWithStats(working_mask, 8)
+    sizes = stats[:, -1][1:]  # Row 0 is background label
+    small_regions = [i + 1 for i, s in enumerate(sizes) if s < area_thresh]
+    if len(small_regions) == 0:
+        return mask, False
+    fill_labels = [0] + small_regions
+    if not correct_holes:
+        fill_labels = [i for i in range(n_labels) if i not in fill_labels]
+        # If every region is below threshold, keep largest
+        if len(fill_labels) == 0:
+            fill_labels = [int(np.argmax(sizes)) + 1]
+    mask = np.isin(regions, fill_labels)
+    return mask, True
+
+
+def coco_encode_rle(uncompressed_rle: Dict[str, Any]) -> Dict[str, Any]:
+    from pycocotools import mask as mask_utils  # type: ignore
+
+    h, w = uncompressed_rle["size"]
+    rle = mask_utils.frPyObjects(uncompressed_rle, h, w)
+    rle["counts"] = rle["counts"].decode("utf-8")  # Necessary to serialize with json
+    return rle
+
+
+def batched_mask_to_box(masks: torch.Tensor) -> torch.Tensor:
+    """
+    Calculates boxes in XYXY format around masks. Return [0,0,0,0] for
+    an empty mask. For input shape C1xC2x...xHxW, the output shape is C1xC2x...x4.
+    """
+    # torch.max below raises an error on empty inputs, just skip in this case
+    if torch.numel(masks) == 0:
+        return torch.zeros(*masks.shape[:-2], 4, device=masks.device)
+
+    # Normalize shape to CxHxW
+    shape = masks.shape
+    h, w = shape[-2:]
+    if len(shape) > 2:
+        masks = masks.flatten(0, -3)
+    else:
+        masks = masks.unsqueeze(0)
+
+    # Get top and bottom edges
+    in_height, _ = torch.max(masks, dim=-1)
+    in_height_coords = in_height * torch.arange(h, device=in_height.device)[None, :]
+    bottom_edges, _ = torch.max(in_height_coords, dim=-1)
+    in_height_coords = in_height_coords + h * (~in_height)
+    top_edges, _ = torch.min(in_height_coords, dim=-1)
+
+    # Get left and right edges
+    in_width, _ = torch.max(masks, dim=-2)
+    in_width_coords = in_width * torch.arange(w, device=in_width.device)[None, :]
+    right_edges, _ = torch.max(in_width_coords, dim=-1)
+    in_width_coords = in_width_coords + w * (~in_width)
+    left_edges, _ = torch.min(in_width_coords, dim=-1)
+
+    # If the mask is empty the right edge will be to the left of the left edge.
+    # Replace these boxes with [0, 0, 0, 0]
+    empty_filter = (right_edges < left_edges) | (bottom_edges < top_edges)
+    out = torch.stack([left_edges, top_edges, right_edges, bottom_edges], dim=-1)
+    out = out * (~empty_filter).unsqueeze(-1)
+
+    # Return to original shape
+    if len(shape) > 2:
+        out = out.reshape(*shape[:-2], 4)
+    else:
+        out = out[0]
+
+    return out

utility/sam/utils/misc.py

@@ -0,0 +1,349 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+import warnings
+from threading import Thread
+
+import numpy as np
+import torch
+from PIL import Image
+from tqdm import tqdm
+
+
+def get_sdpa_settings():
+    if torch.cuda.is_available():
+        old_gpu = torch.cuda.get_device_properties(0).major < 7
+        # only use Flash Attention on Ampere (8.0) or newer GPUs
+        use_flash_attn = torch.cuda.get_device_properties(0).major >= 8
+        if not use_flash_attn:
+            warnings.warn(
+                "Flash Attention is disabled as it requires a GPU with Ampere (8.0) CUDA capability.",
+                category=UserWarning,
+                stacklevel=2,
+            )
+        # keep math kernel for PyTorch versions before 2.2 (Flash Attention v2 is only
+        # available on PyTorch 2.2+, while Flash Attention v1 cannot handle all cases)
+        pytorch_version = tuple(int(v) for v in torch.__version__.split(".")[:2])
+        if pytorch_version < (2, 2):
+            warnings.warn(
+                f"You are using PyTorch {torch.__version__} without Flash Attention v2 support. "
+                "Consider upgrading to PyTorch 2.2+ for Flash Attention v2 (which could be faster).",
+                category=UserWarning,
+                stacklevel=2,
+            )
+        math_kernel_on = pytorch_version < (2, 2) or not use_flash_attn
+    else:
+        old_gpu = True
+        use_flash_attn = False
+        math_kernel_on = True
+
+    return old_gpu, use_flash_attn, math_kernel_on
+
+
+def get_connected_components(mask):
+    """
+    Get the connected components (8-connectivity) of binary masks of shape (N, 1, H, W).
+
+    Inputs:
+    - mask: A binary mask tensor of shape (N, 1, H, W), where 1 is foreground and 0 is
+            background.
+
+    Outputs:
+    - labels: A tensor of shape (N, 1, H, W) containing the connected component labels
+              for foreground pixels and 0 for background pixels.
+    - counts: A tensor of shape (N, 1, H, W) containing the area of the connected
+              components for foreground pixels and 0 for background pixels.
+    """
+    from sam2 import _C
+
+    return _C.get_connected_componnets(mask.to(torch.uint8).contiguous())
+
+
+def mask_to_box(masks: torch.Tensor):
+    """
+    compute bounding box given an input mask
+
+    Inputs:
+    - masks: [B, 1, H, W] masks, dtype=torch.Tensor
+
+    Returns:
+    - box_coords: [B, 1, 4], contains (x, y) coordinates of top left and bottom right box corners, dtype=torch.Tensor
+    """
+    B, _, h, w = masks.shape
+    device = masks.device
+    xs = torch.arange(w, device=device, dtype=torch.int32)
+    ys = torch.arange(h, device=device, dtype=torch.int32)
+    grid_xs, grid_ys = torch.meshgrid(xs, ys, indexing="xy")
+    grid_xs = grid_xs[None, None, ...].expand(B, 1, h, w)
+    grid_ys = grid_ys[None, None, ...].expand(B, 1, h, w)
+    min_xs, _ = torch.min(torch.where(masks, grid_xs, w).flatten(-2), dim=-1)
+    max_xs, _ = torch.max(torch.where(masks, grid_xs, -1).flatten(-2), dim=-1)
+    min_ys, _ = torch.min(torch.where(masks, grid_ys, h).flatten(-2), dim=-1)
+    max_ys, _ = torch.max(torch.where(masks, grid_ys, -1).flatten(-2), dim=-1)
+    bbox_coords = torch.stack((min_xs, min_ys, max_xs, max_ys), dim=-1)
+
+    return bbox_coords
+
+
+def _load_img_as_tensor(img_path, image_size):
+    img_pil = Image.open(img_path)
+    img_np = np.array(img_pil.convert("RGB").resize((image_size, image_size)))
+    if img_np.dtype == np.uint8:  # np.uint8 is expected for JPEG images
+        img_np = img_np / 255.0
+    else:
+        raise RuntimeError(f"Unknown image dtype: {img_np.dtype} on {img_path}")
+    img = torch.from_numpy(img_np).permute(2, 0, 1)
+    video_width, video_height = img_pil.size  # the original video size
+    return img, video_height, video_width
+
+
+class AsyncVideoFrameLoader:
+    """
+    A list of video frames to be load asynchronously without blocking session start.
+    """
+
+    def __init__(
+        self,
+        img_paths,
+        image_size,
+        offload_video_to_cpu,
+        img_mean,
+        img_std,
+        compute_device,
+    ):
+        self.img_paths = img_paths
+        self.image_size = image_size
+        self.offload_video_to_cpu = offload_video_to_cpu
+        self.img_mean = img_mean
+        self.img_std = img_std
+        # items in `self.images` will be loaded asynchronously
+        self.images = [None] * len(img_paths)
+        # catch and raise any exceptions in the async loading thread
+        self.exception = None
+        # video_height and video_width be filled when loading the first image
+        self.video_height = None
+        self.video_width = None
+        self.compute_device = compute_device
+
+        # load the first frame to fill video_height and video_width and also
+        # to cache it (since it's most likely where the user will click)
+        self.__getitem__(0)
+
+        # load the rest of frames asynchronously without blocking the session start
+        def _load_frames():
+            try:
+                for n in tqdm(range(len(self.images)), desc="frame loading"):
+                    self.__getitem__(n)
+            except Exception as e:
+                self.exception = e
+
+        #self.thread = Thread(target=_load_frames, daemon=True)
+        #self.thread.start()
+
+    def __getitem__(self, index):
+        if self.exception is not None:
+            raise RuntimeError("Failure in frame loading thread") from self.exception
+
+        img = self.images[index]
+        if img is not None:
+            return img
+
+        img, video_height, video_width = _load_img_as_tensor(
+            self.img_paths[index], self.image_size
+        )
+        self.video_height = video_height
+        self.video_width = video_width
+        # normalize by mean and std
+        img -= self.img_mean
+        img /= self.img_std
+        if not self.offload_video_to_cpu:
+            img = img.to(self.compute_device, non_blocking=True)
+        #self.images[index] = img
+        return img
+
+    def __len__(self):
+        return len(self.images)
+
+
+def load_video_frames(
+    video_path,
+    image_size,
+    offload_video_to_cpu,
+    img_mean=(0.485, 0.456, 0.406),
+    img_std=(0.229, 0.224, 0.225),
+    async_loading_frames=False,
+    compute_device=torch.device("cuda"),
+):
+    """
+    Load the video frames from video_path. The frames are resized to image_size as in
+    the model and are loaded to GPU if offload_video_to_cpu=False. This is used by the demo.
+    """
+    is_bytes = isinstance(video_path, bytes)
+    is_str = isinstance(video_path, str)
+    is_mp4_path = is_str and os.path.splitext(video_path)[-1] in [".mp4", ".MP4"]
+    if is_bytes or is_mp4_path:
+        return load_video_frames_from_video_file(
+            video_path=video_path,
+            image_size=image_size,
+            offload_video_to_cpu=offload_video_to_cpu,
+            img_mean=img_mean,
+            img_std=img_std,
+            compute_device=compute_device,
+        )
+    elif is_str and os.path.isdir(video_path):
+        return load_video_frames_from_jpg_images(
+            video_path=video_path,
+            image_size=image_size,
+            offload_video_to_cpu=offload_video_to_cpu,
+            img_mean=img_mean,
+            img_std=img_std,
+            async_loading_frames=async_loading_frames,
+            compute_device=compute_device,
+        )
+    else:
+        raise NotImplementedError(
+            "Only MP4 video and JPEG folder are supported at this moment"
+        )
+
+
+def load_video_frames_from_jpg_images(
+    video_path,
+    image_size,
+    offload_video_to_cpu,
+    img_mean=(0.485, 0.456, 0.406),
+    img_std=(0.229, 0.224, 0.225),
+    async_loading_frames=False,
+    compute_device=torch.device("cuda"),
+):
+    """
+    Load the video frames from a directory of JPEG files ("<frame_index>.jpg" format).
+
+    The frames are resized to image_size x image_size and are loaded to GPU if
+    `offload_video_to_cpu` is `False` and to CPU if `offload_video_to_cpu` is `True`.
+
+    You can load a frame asynchronously by setting `async_loading_frames` to `True`.
+    """
+    if isinstance(video_path, str) and os.path.isdir(video_path):
+        jpg_folder = video_path
+    else:
+        raise NotImplementedError(
+            "Only JPEG frames are supported at this moment. For video files, you may use "
+            "ffmpeg (https://ffmpeg.org/) to extract frames into a folder of JPEG files, such as \n"
+            "```\n"
+            "ffmpeg -i <your_video>.mp4 -q:v 2 -start_number 0 <output_dir>/'%05d.jpg'\n"
+            "```\n"
+            "where `-q:v` generates high-quality JPEG frames and `-start_number 0` asks "
+            "ffmpeg to start the JPEG file from 00000.jpg."
+        )
+
+    frame_names = [
+        p
+        for p in os.listdir(jpg_folder)
+        if os.path.splitext(p)[-1] in [".jpg", ".jpeg", ".JPG", ".JPEG", ".png", ".PNG"]
+    ]
+    frame_names.sort(key=lambda p: int(os.path.splitext(p)[0]))
+    num_frames = len(frame_names)
+    if num_frames == 0:
+        raise RuntimeError(f"no images found in {jpg_folder}")
+    img_paths = [os.path.join(jpg_folder, frame_name) for frame_name in frame_names]
+    img_mean = torch.tensor(img_mean, dtype=torch.float32)[:, None, None]
+    img_std = torch.tensor(img_std, dtype=torch.float32)[:, None, None]
+
+    if async_loading_frames:
+        lazy_images = AsyncVideoFrameLoader(
+            img_paths,
+            image_size,
+            offload_video_to_cpu,
+            img_mean,
+            img_std,
+            compute_device,
+        )
+        return lazy_images, lazy_images.video_height, lazy_images.video_width
+
+    images = torch.zeros(num_frames, 3, image_size, image_size, dtype=torch.float32)
+    for n, img_path in enumerate(tqdm(img_paths, desc="frame loading")):
+        images[n], video_height, video_width = _load_img_as_tensor(img_path, image_size)
+    if not offload_video_to_cpu:
+        images = images.to(compute_device)
+        img_mean = img_mean.to(compute_device)
+        img_std = img_std.to(compute_device)
+    # normalize by mean and std
+    images -= img_mean
+    images /= img_std
+    return images, video_height, video_width
+
+
+def load_video_frames_from_video_file(
+    video_path,
+    image_size,
+    offload_video_to_cpu,
+    img_mean=(0.485, 0.456, 0.406),
+    img_std=(0.229, 0.224, 0.225),
+    compute_device=torch.device("cuda"),
+):
+    """Load the video frames from a video file."""
+    import decord
+
+    img_mean = torch.tensor(img_mean, dtype=torch.float32)[:, None, None]
+    img_std = torch.tensor(img_std, dtype=torch.float32)[:, None, None]
+    # Get the original video height and width
+    decord.bridge.set_bridge("torch")
+    video_height, video_width, _ = decord.VideoReader(video_path).next().shape
+    # Iterate over all frames in the video
+    images = []
+    for frame in decord.VideoReader(video_path, width=image_size, height=image_size):
+        images.append(frame.permute(2, 0, 1))
+
+    images = torch.stack(images, dim=0).float() / 255.0
+    if not offload_video_to_cpu:
+        images = images.to(compute_device)
+        img_mean = img_mean.to(compute_device)
+        img_std = img_std.to(compute_device)
+    # normalize by mean and std
+    images -= img_mean
+    images /= img_std
+    return images, video_height, video_width
+
+
+def fill_holes_in_mask_scores(mask, max_area):
+    """
+    A post processor to fill small holes in mask scores with area under `max_area`.
+    """
+    # Holes are those connected components in background with area <= self.max_area
+    # (background regions are those with mask scores <= 0)
+    assert max_area > 0, "max_area must be positive"
+
+    input_mask = mask
+    try:
+        labels, areas = get_connected_components(mask <= 0)
+        is_hole = (labels > 0) & (areas <= max_area)
+        # We fill holes with a small positive mask score (0.1) to change them to foreground.
+        mask = torch.where(is_hole, 0.1, mask)
+    except Exception as e:
+        # Skip the post-processing step on removing small holes if the CUDA kernel fails
+        warnings.warn(
+            f"{e}\n\nSkipping the post-processing step due to the error above. You can "
+            "still use SAM 2 and it's OK to ignore the error above, although some post-processing "
+            "functionality may be limited (which doesn't affect the results in most cases; see "
+            "https://github.com/facebookresearch/sam2/blob/main/INSTALL.md).",
+            category=UserWarning,
+            stacklevel=2,
+        )
+        mask = input_mask
+
+    return mask
+
+
+def concat_points(old_point_inputs, new_points, new_labels):
+    """Add new points and labels to previous point inputs (add at the end)."""
+    if old_point_inputs is None:
+        points, labels = new_points, new_labels
+    else:
+        points = torch.cat([old_point_inputs["point_coords"], new_points], dim=1)
+        labels = torch.cat([old_point_inputs["point_labels"], new_labels], dim=1)
+
+    return {"point_coords": points, "point_labels": labels}

utility/sam/utils/transforms.py

@@ -0,0 +1,118 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import warnings
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.transforms import Normalize, Resize, ToTensor
+
+
+class SAM2Transforms(nn.Module):
+    def __init__(
+        self, resolution, mask_threshold, max_hole_area=0.0, max_sprinkle_area=0.0
+    ):
+        """
+        Transforms for SAM2.
+        """
+        super().__init__()
+        self.resolution = resolution
+        self.mask_threshold = mask_threshold
+        self.max_hole_area = max_hole_area
+        self.max_sprinkle_area = max_sprinkle_area
+        self.mean = [0.485, 0.456, 0.406]
+        self.std = [0.229, 0.224, 0.225]
+        self.to_tensor = ToTensor()
+        self.transforms = torch.jit.script(
+            nn.Sequential(
+                Resize((self.resolution, self.resolution)),
+                Normalize(self.mean, self.std),
+            )
+        )
+
+    def __call__(self, x):
+        x = self.to_tensor(x)
+        return self.transforms(x)
+
+    def forward_batch(self, img_list):
+        img_batch = [self.transforms(self.to_tensor(img)) for img in img_list]
+        img_batch = torch.stack(img_batch, dim=0)
+        return img_batch
+
+    def transform_coords(
+        self, coords: torch.Tensor, normalize=False, orig_hw=None
+    ) -> torch.Tensor:
+        """
+        Expects a torch tensor with length 2 in the last dimension. The coordinates can be in absolute image or normalized coordinates,
+        If the coords are in absolute image coordinates, normalize should be set to True and original image size is required.
+
+        Returns
+            Un-normalized coordinates in the range of [0, 1] which is expected by the SAM2 model.
+        """
+        if normalize:
+            assert orig_hw is not None
+            h, w = orig_hw
+            coords = coords.clone()
+            coords[..., 0] = coords[..., 0] / w
+            coords[..., 1] = coords[..., 1] / h
+
+        coords = coords * self.resolution  # unnormalize coords
+        return coords
+
+    def transform_boxes(
+        self, boxes: torch.Tensor, normalize=False, orig_hw=None
+    ) -> torch.Tensor:
+        """
+        Expects a tensor of shape Bx4. The coordinates can be in absolute image or normalized coordinates,
+        if the coords are in absolute image coordinates, normalize should be set to True and original image size is required.
+        """
+        boxes = self.transform_coords(boxes.reshape(-1, 2, 2), normalize, orig_hw)
+        return boxes
+
+    def postprocess_masks(self, masks: torch.Tensor, orig_hw) -> torch.Tensor:
+        """
+        Perform PostProcessing on output masks.
+        """
+        from sam2.utils.misc import get_connected_components
+
+        masks = masks.float()
+        input_masks = masks
+        mask_flat = masks.flatten(0, 1).unsqueeze(1)  # flatten as 1-channel image
+        try:
+            if self.max_hole_area > 0:
+                # Holes are those connected components in background with area <= self.fill_hole_area
+                # (background regions are those with mask scores <= self.mask_threshold)
+                labels, areas = get_connected_components(
+                    mask_flat <= self.mask_threshold
+                )
+                is_hole = (labels > 0) & (areas <= self.max_hole_area)
+                is_hole = is_hole.reshape_as(masks)
+                # We fill holes with a small positive mask score (10.0) to change them to foreground.
+                masks = torch.where(is_hole, self.mask_threshold + 10.0, masks)
+
+            if self.max_sprinkle_area > 0:
+                labels, areas = get_connected_components(
+                    mask_flat > self.mask_threshold
+                )
+                is_hole = (labels > 0) & (areas <= self.max_sprinkle_area)
+                is_hole = is_hole.reshape_as(masks)
+                # We fill holes with negative mask score (-10.0) to change them to background.
+                masks = torch.where(is_hole, self.mask_threshold - 10.0, masks)
+        except Exception as e:
+            # Skip the post-processing step if the CUDA kernel fails
+            warnings.warn(
+                f"{e}\n\nSkipping the post-processing step due to the error above. You can "
+                "still use SAM 2 and it's OK to ignore the error above, although some post-processing "
+                "functionality may be limited (which doesn't affect the results in most cases; see "
+                "https://github.com/facebookresearch/sam2/blob/main/INSTALL.md).",
+                category=UserWarning,
+                stacklevel=2,
+            )
+            masks = input_masks
+
+        masks = F.interpolate(masks, orig_hw, mode="bilinear", align_corners=False)
+        return masks

web/power_lora_loader_js/widgets.js

@@ -289,7 +289,7 @@ export class StrengthCopyWidget extends RgthreeBaseWidget {
         // Draw label and toggle from right to left
         const toggleBounds = drawTogglePart(ctx, [posX - height * 1.5, posY], height, outputBothValue);
         this.hitAreas['output_both_toggle'].bounds = toggleBounds;
-        const labelText = "Use loras on both";
+        const labelText = "Output both";
         ctx.fillText(labelText, posX - height * 1.5 - ctx.measureText(labelText).width - innerMargin, midY);
 
         // Adjust rightX for buttons (they stay in same position)

web/power_spline_editor/canvas/canvas_box_timeline.js

@@ -413,4 +413,39 @@ export function attachBoxTimelineHelpers(editor) {
       };
     }
   };
+
+  editor.moveBoxLayerKey = (widget, fromFrame, toFrame) => {
+    if (!editor._isBoxLayerWidget(widget)) return false;
+    const keys = editor._ensureBoxLayerData(widget) || [];
+    const maxFrames = editor._getMaxFrames();
+    const fromFrameClamped = Math.max(1, Math.min(maxFrames, Math.round(fromFrame || 1)));
+    const toFrameClamped = Math.max(1, Math.min(maxFrames, Math.round(toFrame || 1)));
+
+    if (fromFrameClamped === toFrameClamped) return false;
+
+    const fromIdx = keys.findIndex(k => k.frame === fromFrameClamped);
+    if (fromIdx === -1) return false; // No keyframe at source frame
+
+    const keyToMove = { ...keys[fromIdx], frame: toFrameClamped };
+
+    // Remove keyframe at target frame if it exists (overwrite behavior)
+    const toIdx = keys.findIndex(k => k.frame === toFrameClamped);
+    if (toIdx >= 0) {
+      keys.splice(toIdx, 1);
+    }
+
+    // Remove from old position and add at new position
+    keys.splice(fromIdx, 1);
+    keys.push(keyToMove);
+    keys.sort((a, b) => a.frame - b.frame);
+
+    widget.value.box_keys = keys;
+    widget.value.box_timeline_point = toFrameClamped;
+    editor.applyBoxTimelineFrame(widget, toFrameClamped);
+    editor._forceRebuildNextRender = true;
+    try { editor.layerRenderer?.render(); } catch {}
+    try { editor.node?.setDirtyCanvas?.(true, true); } catch {}
+
+    return true;
+  };
 }

web/power_spline_editor/canvas/canvas_main.js

@@ -94,7 +94,7 @@ export default class SplineEditor2 {
     // Global toggle for inactive flow animation (default OFF)
     this._inactiveFlowEnabled = false;
 
-    this._shortcutKeys = { s: 0 }; // Initialize 's' key state with a timestamp (0 means not pressed)
+    this._shortcutKeys = { s: 0, d: 0 }; // Initialize 's' and 'd' key state with a timestamp (0 means not pressed)
     this._boxPlayModeStartedAt = 0;
     this._boxRefOnlyMode = false;
     this._boxPlayStopGuardMs = 120;
@@ -157,6 +157,9 @@ export default class SplineEditor2 {
       if (key === 's') {
         this._shortcutKeys.s = performance.now(); // Store timestamp of keydown
       }
+      if (key === 'd') {
+        this._shortcutKeys.d = performance.now(); // Store timestamp of keydown
+      }
     };
     const handleShortcutKeyUp = (ev) => {
       const key = ev?.key?.toLowerCase?.();
@@ -164,6 +167,9 @@ export default class SplineEditor2 {
       if (key === 's') {
         this._shortcutKeys.s = 0; // Reset to 0 on keyup
       }
+      if (key === 'd') {
+        this._shortcutKeys.d = 0; // Reset to 0 on keyup
+      }
       if (key === ' ' && !ev.shiftKey && this._boxPlayModeActive) {
         const elapsed = performance.now() - this._boxPlayModeStartedAt;
         if (elapsed >= this._boxPlayStopGuardMs) {

web/power_spline_editor/canvas/canvas_sam2.js

@@ -0,0 +1,435 @@
+/**
+ * SAM2 Point-Click Masking Helpers
+ *
+ * This module provides SAM2 (Segment Anything Model 2) point-click masking
+ * functionality for the RefCanvasEditor. It allows users to create masks by
+ * clicking on objects instead of drawing them manually.
+ *
+ * Main features:
+ * - Mode management (enter/exit SAM2 mode)
+ * - Point management (add/remove SAM2 click points)
+ * - API integration (debounced prediction requests)
+ * - Visual feedback (points rendering and mask preview)
+ * - Auto-commit on mode switch
+ */
+
+/**
+ * Attach SAM2 helpers to a RefCanvasEditor instance
+ * @param {Object} editor - The RefCanvasEditor instance to extend
+ */
+export function attachSam2Helpers(editor) {
+  // SAM2 mode state
+  editor._sam2Mode = false;
+  editor._sam2Points = [];
+  editor._sam2MaskPath = null;
+  editor._sam2DebounceTimer = null;
+  editor._sam2RenderCallbacks = []; // Store render callbacks
+
+  /**
+   * Toggle SAM2 mode on/off
+   */
+  editor.toggleSam2Mode = () => {
+    if (editor._sam2Mode) {
+      editor.exitSam2Mode();
+    } else {
+      editor.enterSam2Mode();
+    }
+  };
+
+  /**
+   * Enter SAM2 mode - exit pencil mode first
+   */
+  editor.enterSam2Mode = () => {
+    // Exit pencil mode if active
+    if (editor._lassoDrawingActive) {
+      editor.exitLassoMode?.();
+      // Clear pencil masks
+      editor.clearLayerMasks?.();
+    }
+
+    editor._sam2Mode = true;
+    editor._sam2Points = [];
+    editor._sam2MaskPath = null;
+
+    // Set the active layer (same as lasso mode)
+    editor._sam2ActiveLayer = editor.selectedRefLayer;
+
+    editor.updateMagicWandButton(true);
+    editor.forceCanvasRefresh?.();
+  };
+
+  /**
+   * Exit SAM2 mode - bake mask to pencil mode and clean up
+   */
+  editor.exitSam2Mode = () => {
+    // Step 1: Bake SAM2 mask to pencil mode (convert to lasso format)
+    if (editor._sam2MaskPath && editor._sam2Points.length > 0) {
+      editor.bakeSam2MaskToLasso();
+    }
+
+    // Step 2: Update wand button to deselected state (BEFORE clearing layer reference)
+    editor.updateMagicWandButton(false);
+
+    // Step 3: Clear all SAM2 state
+    editor._sam2Mode = false;
+    editor._sam2ActiveLayer = null;
+    editor._sam2Points = [];
+    editor._sam2MaskPath = null;
+
+    // Step 4: Clear SAM2 visuals from canvas
+    editor.clearSam2Visuals();
+    editor.forceCanvasRefresh?.();
+  };
+
+  /**
+   * Update magic wand button visual state
+   * @param {boolean} active - Whether the button should be active
+   */
+  editor.updateMagicWandButton = (active) => {
+    if (!editor._sam2ActiveLayer || !editor._sam2ActiveLayer.magicWandButton) return;
+
+    const btn = editor._sam2ActiveLayer.magicWandButton;
+    if (btn) {
+      if (active) {
+        btn.classList.add('active');
+        btn.style.background = '#4CAF50';
+        btn.style.color = 'white';
+      } else {
+        btn.classList.remove('active');
+        btn.style.background = 'transparent';
+        btn.style.color = '#888';
+      }
+    }
+  };
+
+  /**
+   * Add SAM2 point and trigger prediction
+   * @param {number} x - X coordinate in video space
+   * @param {number} y - Y coordinate in video space
+   */
+  editor.addSam2Point = (x, y) => {
+    const point = {
+      x: x,
+      y: y,
+      uid: Date.now() + Math.random()
+    };
+    editor._sam2Points.push(point);
+
+    // Render immediately to show the point
+    editor.renderSam2PointsImmediate();
+
+    editor.requestSam2Prediction();
+  };
+
+  /**
+   * Render SAM2 points immediately (for instant visual feedback)
+   */
+  editor.renderSam2PointsImmediate = () => {
+    if (!editor.refCanvasEditor || !editor.refCanvasEditor.ctx) return;
+
+    const ctx = editor.refCanvasEditor.ctx;
+    const lastPoint = editor._sam2Points[editor._sam2Points.length - 1];
+    if (!lastPoint) return;
+
+    // SAM2 points are stored in video space, need to convert to canvas space
+    const canvasPos = editor.refCanvasEditor.transformToCanvasSpace?.(lastPoint.x, lastPoint.y) || lastPoint;
+    const index = editor._sam2Points.length - 1;
+
+    ctx.save();
+
+    // Draw red circle with white border
+    ctx.beginPath();
+    ctx.arc(canvasPos.x, canvasPos.y, 8, 0, Math.PI * 2);
+    ctx.fillStyle = '#f44336';  // Red
+    ctx.fill();
+    ctx.strokeStyle = 'white';
+    ctx.lineWidth = 3;
+    ctx.stroke();
+
+    // Draw point number
+    ctx.fillStyle = 'white';
+    ctx.font = 'bold 10px Arial';
+    ctx.textAlign = 'center';
+    ctx.textBaseline = 'middle';
+    ctx.fillText(index + 1, canvasPos.x, canvasPos.y);
+
+    ctx.restore();
+  };
+
+  /**
+   * Remove SAM2 point by UID
+   * @param {number} uid - Unique identifier of the point to remove
+   */
+  editor.removeSam2Point = (uid) => {
+    editor._sam2Points = editor._sam2Points.filter(p => p.uid !== uid);
+    editor.refreshCanvas?.();
+
+    // Rerun prediction if there are still points left
+    if (editor._sam2Points.length > 0) {
+      editor.requestSam2Prediction();
+    } else {
+      // Clear mask if no points left
+      editor.clearSam2MaskPreview();
+    }
+  };
+
+  /**
+   * Find SAM2 point at canvas coordinates
+   * @param {number} canvasX - X coordinate in canvas space
+   * @param {number} canvasY - Y coordinate in canvas space
+   * @param {number} threshold - Distance threshold for finding points (default: 10)
+   * @returns {Object|null} The found point or null
+   */
+  editor.findSam2PointAt = (canvasX, canvasY, threshold = 10) => {
+    if (!editor.refCanvasEditor) return null;
+
+    for (const point of editor._sam2Points) {
+      // SAM2 points are stored in video space, need to convert to canvas space
+      const canvasPos = editor.refCanvasEditor.transformToCanvasSpace?.(point.x, point.y) || point;
+
+      const dist = Math.sqrt(
+        Math.pow(canvasPos.x - canvasX, 2) +
+        Math.pow(canvasPos.y - canvasY, 2)
+      );
+      if (dist < threshold) {
+        return point;
+      }
+    }
+    return null;
+  };
+
+  /**
+   * Request SAM2 prediction (debounced)
+   * Debounces API calls to avoid excessive requests
+   */
+  editor.requestSam2Prediction = () => {
+    // Clear existing timer
+    if (editor._sam2DebounceTimer) {
+      clearTimeout(editor._sam2DebounceTimer);
+    }
+
+    // Debounce API call
+    editor._sam2DebounceTimer = setTimeout(async () => {
+      if (editor._sam2Points.length === 0) {
+        editor.clearSam2MaskPreview();
+        return;
+      }
+
+      await editor.performSam2Prediction();
+    }, 500); // 500ms debounce
+  };
+
+  /**
+   * Perform SAM2 prediction via API
+   */
+  editor.performSam2Prediction = async () => {
+    try {
+      // Get current canvas image as base64
+      const imageData = editor.getCanvasImageAsBase64();
+      if (!imageData) {
+        console.error("[SAM2] Failed to get canvas image");
+        return;
+      }
+
+      // Prepare API request
+      const response = await fetch('/wanvideowrapper_qq/sam2_predict', {
+        method: 'POST',
+        headers: { 'Content-Type': 'application/json' },
+        body: JSON.stringify({
+          image: imageData,
+          points: editor._sam2Points.map(p => ({
+            x: p.x,
+            y: p.y,
+            label: 1  // All points are foreground for now
+          }))
+        })
+      });
+
+      if (!response.ok) {
+        const error = await response.json();
+        console.error("[SAM2] API error:", error);
+        return;
+      }
+
+      const result = await response.json();
+      if (result.success && result.path) {
+        editor._sam2MaskPath = result.path;
+        editor.renderSam2MaskPreview(result.path);
+        editor.refreshCanvas?.(); // Trigger canvas refresh to show the mask
+      }
+
+    } catch (error) {
+      console.error("[SAM2] Prediction failed:", error);
+    }
+  };
+
+  /**
+   * Clear SAM2 visuals (points and mask preview)
+   */
+  editor.clearSam2Visuals = () => {
+    // Visuals will be cleared naturally by not rendering them
+  };
+
+  /**
+   * Render SAM2 mask preview
+   * @param {Array} path - Array of {x, y} normalized coordinates
+   */
+  editor.renderSam2MaskPreview = (path) => {
+    if (!editor.refCanvasEditor) return;
+
+    const ctx = editor.refCanvasEditor.ctx;
+    if (!ctx || !path || path.length === 0) return;
+
+    // Store the path for rendering
+    editor._sam2MaskPath = path;
+    editor.forceCanvasRefresh?.();
+  };
+
+  /**
+   * Render SAM2 visuals (points and mask preview) on canvas
+   * This function is called during canvas refresh
+   * @param {CanvasRenderingContext2D} ctx - Canvas context
+   * @param {number} width - Canvas width
+   * @param {number} height - Canvas height
+   */
+  editor.renderSam2Visuals = (ctx, width, height) => {
+    if (!editor._sam2Mode) return;
+
+    // Render mask preview if available
+    if (editor._sam2MaskPath && editor._sam2MaskPath.length > 0) {
+      ctx.save();
+
+      // Convert normalized path to canvas coordinates
+      const canvasPath = editor._sam2MaskPath.map(pt => {
+        // Path points are normalized (0-1), need to convert to canvas space
+        const canvasPos = editor.refCanvasEditor.denormalizePoint(pt);
+        return { x: canvasPos.x, y: canvasPos.y };
+      });
+
+      // Draw filled path with green transparency
+      if (canvasPath.length > 2) {
+        ctx.beginPath();
+        ctx.moveTo(canvasPath[0].x, canvasPath[0].y);
+        for (let i = 1; i < canvasPath.length; i++) {
+          ctx.lineTo(canvasPath[i].x, canvasPath[i].y);
+        }
+        ctx.closePath();
+
+        // Fill with green transparency
+        ctx.fillStyle = 'rgba(76, 175, 80, 0.3)';
+        ctx.fill();
+
+        // Draw dashed green border
+        ctx.strokeStyle = '#4CAF50';
+        ctx.lineWidth = 2;
+        ctx.setLineDash([5, 5]); // Dashed line
+        ctx.stroke();
+      }
+
+      ctx.restore();
+    }
+
+    // Render SAM2 points
+    if (editor._sam2Points.length > 0) {
+      ctx.save();
+
+      editor._sam2Points.forEach((point, index) => {
+        // SAM2 points are stored in video space, need to convert to canvas space
+        const canvasPos = editor.refCanvasEditor.transformToCanvasSpace?.(point.x, point.y) || point;
+
+        // Draw red circle with white border (larger for visibility)
+        ctx.beginPath();
+        ctx.arc(canvasPos.x, canvasPos.y, 8, 0, Math.PI * 2);
+        ctx.fillStyle = '#f44336';  // Red
+        ctx.fill();
+        ctx.strokeStyle = 'white';
+        ctx.lineWidth = 3;
+        ctx.stroke();
+
+        // Draw point number for better visibility
+        ctx.fillStyle = 'white';
+        ctx.font = 'bold 10px Arial';
+        ctx.textAlign = 'center';
+        ctx.textBaseline = 'middle';
+        ctx.fillText(index + 1, canvasPos.x, canvasPos.y);
+      });
+
+      ctx.restore();
+    }
+  };
+
+  /**
+   * Clear SAM2 mask preview only
+   */
+  editor.clearSam2MaskPreview = () => {
+    editor._sam2MaskPath = null;
+    editor.refreshCanvas?.(); // Trigger canvas refresh to clear the mask
+  };
+
+  /**
+   * Get current canvas image as base64
+   * @returns {string|null} Base64 encoded image or null
+   */
+  editor.getCanvasImageAsBase64 = () => {
+    try {
+      // Get the background image from the canvas editor
+      if (editor.refCanvasEditor && editor.refCanvasEditor.backgroundImage) {
+        // Create canvas to draw image
+        const canvas = document.createElement('canvas');
+        canvas.width = editor.refCanvasEditor.originalImageWidth || editor.refCanvasEditor.width;
+        canvas.height = editor.refCanvasEditor.originalImageHeight || editor.refCanvasEditor.height;
+        const ctx = canvas.getContext('2d');
+
+        // Draw background image at original resolution
+        ctx.drawImage(editor.refCanvasEditor.backgroundImage, 0, 0, canvas.width, canvas.height);
+
+        // Convert to RGB by creating a temporary canvas and drawing without alpha
+        const tempCanvas = document.createElement('canvas');
+        tempCanvas.width = canvas.width;
+        tempCanvas.height = canvas.height;
+        const tempCtx = tempCanvas.getContext('2d');
+
+        // Draw the image (this automatically handles RGBA to RGB conversion)
+        tempCtx.drawImage(canvas, 0, 0);
+
+        // Return as base64 JPEG (forces RGB conversion, no alpha channel)
+        return tempCanvas.toDataURL('image/jpeg', 0.95);
+      }
+      return null;
+    } catch (error) {
+      console.error("[SAM2] Failed to get canvas image:", error);
+      return null;
+    }
+  };
+
+  /**
+   * Bake SAM2 mask to lasso format for pencil mode editing
+   * Converts the current SAM2 mask to lasso format and adds it to the layer
+   */
+  editor.bakeSam2MaskToLasso = () => {
+    if (!editor._sam2MaskPath || editor._sam2MaskPath.length === 0) {
+      return;
+    }
+
+    if (!editor._sam2ActiveLayer) {
+      return;
+    }
+
+    // Get current layer data
+    const layerData = editor._sam2ActiveLayer.value;
+
+    // Initialize lassoShape if needed
+    if (!layerData.lassoShape) {
+      layerData.lassoShape = {
+        additivePaths: [],
+        subtractivePaths: []
+      };
+    }
+
+    // Add SAM2 mask as additive path
+    layerData.lassoShape.additivePaths.push(editor._sam2MaskPath);
+
+    // Mark that this layer now has lasso data for pencil mode
+    editor._sam2ActiveLayer.hasLassoData = true;
+  };
+}

web/power_spline_editor/canvas_lasso.js

@@ -517,9 +517,39 @@ function exitLassoMode(node) {
 }
 
 /**
- * Handle canvas mousedown - start drawing
+ * Handle canvas mousedown - start drawing or handle SAM2 mode
  */
 function handleCanvasMouseDown(node, e) {
+  // Handle SAM2 mode if active
+  if (node._sam2Mode && e.button === 0) {
+    const coords = node.refCanvasEditor.getCanvasCoords(e);
+    const mouseX = coords.x;
+    const mouseY = coords.y;
+
+    // Ensure coordinates are within the canvas bounds
+    if (mouseX < 0 || mouseX > node.refsCanvas.width || mouseY < 0 || mouseY > node.refsCanvas.height) {
+      return; // Ignore events outside the canvas
+    }
+
+    // Check if clicking on existing point AND holding Ctrl/meta key
+    if ((e.ctrlKey || e.metaKey)) {
+      const clickedPoint = node.findSam2PointAt?.(mouseX, mouseY);
+      if (clickedPoint) {
+        // Ctrl+click on existing point - remove it and rerun prediction
+        node.removeSam2Point?.(clickedPoint.uid);
+        return;
+      }
+      // Ctrl+click on empty space - still add point (allows adding points with Ctrl held)
+    }
+
+    // Add new point (regular click, Shift+click, or Ctrl+click on empty space)
+    // All these cases add points
+    const videoSpace = node.refCanvasEditor.transformToVideoSpace(mouseX, mouseY);
+    node.addSam2Point?.(videoSpace.x, videoSpace.y);
+    return;
+  }
+
+  // Original lasso handling
   if (!node._lassoDrawingActive || !node._lassoActiveLayer) return;
 
   // Use RefCanvas to get canvas coordinates (same as PowerSplineEditor)

web/power_spline_editor/canvas_main_ref.js

@@ -155,4 +155,52 @@ export default class RefCanvas {
       y: (event.clientY - rect.top) * scaleY
     };
   }
+
+  /**
+   * Denormalize a single point - converts 0-1 normalized to canvas coordinates
+   * @param {Object} point - Point with {x, y} in normalized 0-1 coordinates
+   * @returns {Object} Point with {x, y} in canvas coordinates
+   */
+  denormalizePoint(point) {
+    const result = this.denormalizePoints([point]);
+    return result[0] || point;
+  }
+
+  /**
+   * Transform canvas coordinates to video space (original image pixel coordinates)
+   * @param {number} canvasX - X coordinate in canvas space
+   * @param {number} canvasY - Y coordinate in canvas space
+   * @returns {Object} {x, y} in video space (original image pixel coordinates)
+   */
+  transformToVideoSpace(canvasX, canvasY) {
+    if (!this.originalImageWidth || !this.originalImageHeight || this.scale <= 0) {
+      // Fallback: return canvas coordinates if no transformation available
+      return { x: canvasX, y: canvasY };
+    }
+
+    // Reverse the denormalization: canvas -> original image coordinates
+    const videoX = (canvasX - this.offsetX) / this.scale;
+    const videoY = (canvasY - this.offsetY) / this.scale;
+
+    return { x: videoX, y: videoY };
+  }
+
+  /**
+   * Transform video space coordinates to canvas coordinates
+   * @param {number} videoX - X coordinate in video space (original image pixels)
+   * @param {number} videoY - Y coordinate in video space (original image pixels)
+   * @returns {Object} {x, y} in canvas coordinates
+   */
+  transformToCanvasSpace(videoX, videoY) {
+    if (!this.originalImageWidth || !this.originalImageHeight || this.scale <= 0) {
+      // Fallback: return video coordinates if no transformation available
+      return { x: videoX, y: videoY };
+    }
+
+    // Apply the denormalization: video space -> canvas coordinates
+    const canvasX = (videoX * this.scale) + this.offsetX;
+    const canvasY = (videoY * this.scale) + this.offsetY;
+
+    return { x: canvasX, y: canvasY };
+  }
 }

web/power_spline_editor/canvas_top_row.js

@@ -174,11 +174,10 @@ export async function handlePrepareRefsRefresh(node) {
   console.log("[handlePrepareRefsRefresh] Refresh button clicked");
 
   try {
-    // Try to get image from connected node (ImageResizeKJv2 or LoadImage)
-    const { findConnectedSourceNode, getReferenceImageFromConnectedNode } =
-      await import('./graph_query.js');
+    const { findConnectedSourceNode, getReferenceImageFromConnectedNode } = await import('./graph_query.js');
 
     let imageData = null;
+    let isVideo = false;
 
     // Check for connected source on bg_image input
     const sourceNodeObj = findConnectedSourceNode(node, 'bg_image');
@@ -186,65 +185,184 @@ export async function handlePrepareRefsRefresh(node) {
     if (sourceNodeObj?.node) {
       console.log("[handlePrepareRefsRefresh] Found connected node:", sourceNodeObj.node.type);
 
-      // Get image from connected node
-      imageData = await getReferenceImageFromConnectedNode(node, 'bg_image');
+      let sourceNode = sourceNodeObj.node;
+      let resizeParams = null;
 
-      if (imageData) {
-        console.log("[handlePrepareRefsRefresh] Successfully loaded image from connected node");
+      // Check if the connected node is ImageResizeKJv2, and if so, follow the chain to find the real source
+      if (sourceNode.type === 'ImageResizeKJv2') {
+        console.log("[handlePrepareRefsRefresh] Found ImageResizeKJv2, following chain to find video source");
 
-        // Load the image to the canvas
-        const img = await new Promise((resolve, reject) => {
-          const image = new Image();
-          image.onload = () => {
-            node.refCanvasEditor.loadBackgroundImage(image);
-            resolve(image);
-          };
-          image.onerror = reject;
-          image.src = imageData;
-        });
-
-        // Store for persistence
-        node.loadedBgImage = img;
-
-        // Save via backend
-        try {
-          await fetch('/wanvideowrapper_qq/save_prepare_refs_images', {
-            method: 'POST',
-            headers: { 'Content-Type': 'application/json' },
-            body: JSON.stringify({ bg_image: imageData, ref_images: [] })
-          })
-          .then(r => r.json())
-          .then(res => {
-            if (res.success && res.paths?.bg_image_path) {
-              node.properties.bg_image_path = res.paths.bg_image_path;
-              try {
-                sessionStorage.setItem(
-                  `prepare-refs-bg-${node.uuid}`,
-                  res.paths.bg_image_path
-                );
-              } catch (e) { /* ignore */ }
-            }
-          })
-          .catch(console.error);
-        } catch (e) {
-          console.warn('[handlePrepareRefsRefresh] Failed to save image to backend:', e);
+        // Extract resize parameters from ImageResizeKJv2
+        resizeParams = {};
+        if (sourceNode.widgets) {
+          for (const widget of sourceNode.widgets) {
+            resizeParams[widget.name] = widget.value;
+          }
         }
+        console.log("[handlePrepareRefsRefresh] Extracted resize params:", resizeParams);
 
-        // Refresh canvas display
-        node.forceCanvasRefresh();
-        return;
+        // Find the actual source connected to ImageResizeKJv2
+        const actualSourceObj = findConnectedSourceNode(sourceNode, 'image');
+        if (actualSourceObj?.node) {
+          sourceNode = actualSourceObj.node;
+          console.log("[handlePrepareRefsRefresh] Found actual source after ImageResizeKJv2:", sourceNode.type);
+        } else {
+          console.log("[handlePrepareRefsRefresh] No source found after ImageResizeKJv2");
+          return;
+        }
       }
+
+      // Check common video node types
+      const videoNodeTypes = ['VHS_LoadVideo', 'VHS_LoadVideoUpload', 'Load Video Path', 'VideoImporter'];
+      if (videoNodeTypes.some(type => sourceNode.type.includes(type))) {
+        isVideo = true;
+        console.log("[handlePrepareRefsRefresh] Detected video node type:", sourceNode.type);
+      }
+
+      // For video nodes, extract first frame using HTML5 video element
+      if (isVideo) {
+        console.log("[handlePrepareRefsRefresh] Handling video node - extracting first frame");
+
+        try {
+          // Get the video path from VHS node widgets
+          const videoWidget = sourceNode.widgets?.find(w => w.name === 'video' || w.name === 'videoprefix');
+
+          if (videoWidget && videoWidget.value) {
+            const videoPath = videoWidget.value;
+            console.log("[handlePrepareRefsRefresh] Found video path:", videoPath);
+
+            // Build video URL using ComfyUI's /view endpoint
+            const videoUrl = `/view?filename=${encodeURIComponent(videoPath)}&subfolder=&type=input`;
+            console.log("[handlePrepareRefsRefresh] Loading video from:", videoUrl);
+
+            // Create a temporary video element to extract the first frame
+            const videoElement = document.createElement('video');
+            videoElement.muted = true;
+            videoElement.playsInline = true;
+            videoElement.preload = 'auto';
+            videoElement.crossOrigin = 'anonymous';
+
+            // Wait for video to load first frame
+            imageData = await new Promise((resolve, reject) => {
+              const onLoaded = () => {
+                console.log("[handlePrepareRefsRefresh] Video loaded, extracting first frame");
+
+                try {
+                  // Create canvas to capture frame
+                  let canvas = document.createElement('canvas');
+                  canvas.width = videoElement.videoWidth;
+                  canvas.height = videoElement.videoHeight;
+                  const ctx = canvas.getContext('2d');
+
+                  // Draw first frame
+                  ctx.drawImage(videoElement, 0, 0);
+
+                  // Apply resize if ImageResizeKJv2 detected
+                  if (resizeParams) {
+                    console.log("[handlePrepareRefsRefresh] Applying ImageResizeKJv2 parameters");
+
+                    // Convert canvas to image for resizing
+                    const img = new Image();
+                    img.onload = async () => {
+                      try {
+                        const { applyResizeParams } = await import('./image_resize.js');
+                        const result = await applyResizeParams(img, resizeParams);
+
+                        // Update canvas with resized image
+                        canvas = result.canvas;
+
+                        // Convert to base64
+                        const dataUrl = canvas.toDataURL('image/png');
+                        resolve(dataUrl);
+                      } catch (e) {
+                        console.error("[handlePrepareRefsRefresh] Error applying resize:", e);
+                        reject(e);
+                      } finally {
+                        // Cleanup
+                        URL.revokeObjectURL(img.src);
+                      }
+                    };
+                    img.onerror = () => reject(new Error("Failed to load canvas as image"));
+
+                    // Convert canvas to blob for the image
+                    canvas.toBlob((blob) => {
+                      const url = URL.createObjectURL(blob);
+                      img.src = url;
+                    }, 'image/png');
+                  } else {
+                    // No resize, just convert canvas to base64
+                    const dataUrl = canvas.toDataURL('image/png');
+                    resolve(dataUrl);
+                  }
+                } catch (e) {
+                  console.error("[handlePrepareRefsRefresh] Error extracting frame:", e);
+                  reject(e);
+                } finally {
+                  // Cleanup
+                  videoElement.removeEventListener('loadeddata', onLoaded);
+                  videoElement.removeEventListener('error', onError);
+                  videoElement.remove();
+                }
+              };
+
+              const onError = (e) => {
+                console.error("[handlePrepareRefsRefresh] Video loading error:", e);
+                reject(new Error("Failed to load video"));
+                videoElement.removeEventListener('loadeddata', onLoaded);
+                videoElement.removeEventListener('error', onError);
+              };
+
+              videoElement.addEventListener('loadeddata', onLoaded);
+              videoElement.addEventListener('error', onError);
+
+              // Set source to start loading
+              videoElement.src = videoUrl;
+
+              // Set a timeout in case video doesn't load
+              setTimeout(() => {
+                if (videoElement.readyState < 2) {
+                  reject(new Error("Video loading timeout"));
+                }
+              }, 10000); // 10 second timeout
+            });
+
+            console.log("[handlePrepareRefsRefresh] Successfully extracted first frame from video");
+          } else {
+            console.log("[handlePrepareRefsRefresh] No video widget found in node");
+          }
+        } catch (e) {
+          console.error("[handlePrepareRefsRefresh] Error handling video:", e);
+          throw e;
+        }
+      } else {
+        // Handle image nodes
+        console.log("[handlePrepareRefsRefresh] Handling image node");
+        imageData = await getReferenceImageFromConnectedNode(sourceNode);
+      }
+    } else {
+      console.log("[handlePrepareRefsRefresh] No connected source node found");
     }
 
-    console.log("[handlePrepareRefsRefresh] No connected image node found");
+    // If we got image data, load it into the Prepare Refs canvas
+    if (imageData) {
+      console.log("[handlePrepareRefsRefresh] Loading image into Prepare Refs canvas");
 
-  } catch (e) {
-    console.error('[handlePrepareRefsRefresh] Error:', e);
-  }
+      const img = new Image();
+      img.onload = () => {
+        if (node.refCanvasEditor) {
+          node.refCanvasEditor.loadBackgroundImage(img);
+          node.refreshCanvas();
+          node.forceCanvasRefresh?.();
+        }
+      };
+      img.onerror = () => console.error("[handlePrepareRefsRefresh] Failed to load extracted image");
+      img.src = imageData;
+    } else {
+      console.log("[handlePrepareRefsRefresh] No image data obtained");
+    }
 
-  // Always force refresh at the end
-  if (node.forceCanvasRefresh) {
-    node.forceCanvasRefresh();
+  } catch (error) {
+    console.error("[handlePrepareRefsRefresh] Error during refresh:", error);
   }
 }
 

web/power_spline_editor/layer_type_box.js

@@ -66,6 +66,9 @@ export class BoxLayerWidget extends RgthreeBaseWidget {
         this._buttonPressDurationMs = 140;
         this._playbackInterval = null;
         this._timelinePlayIcons = ['▶', '❚❚']; // fixed thick set
+        // Keyframe drag state
+        this._keyframeDragActive = false;
+        this._keyframeDragStartFrame = null;
     }
 
     _getMaxFrames(node) {
@@ -343,21 +346,36 @@ export class BoxLayerWidget extends RgthreeBaseWidget {
             : [];
         if (boxKeys.length) {
             ctx.save();
-            ctx.fillStyle = '#ff8c00';
-            ctx.strokeStyle = '#ffffff';
             const denom = sliderSteps > 1 ? (sliderSteps - 1) : 1;
             const currentFrame = sliderValue;
             for (const key of boxKeys) {
                 const keyFrame = Math.max(1, Math.min(sliderSteps, Math.round(key.frame || 1)));
                 const progress = (keyFrame - 1) / denom;
                 const keyX = timelineX + progress * timelineWidth;
-                const radius = currentFrame === keyFrame ? 5 : 4;
-                ctx.beginPath();
-                ctx.arc(keyX, timelineMidY, radius, 0, Math.PI * 2);
-                ctx.fill();
-                if (currentFrame === keyFrame) {
-                    ctx.lineWidth = 1.5;
+
+                // Highlight keyframe being dragged
+                const isDraggingKeyframe = this._keyframeDragActive && this._keyframeDragStartFrame === keyFrame;
+
+                if (isDraggingKeyframe) {
+                    ctx.fillStyle = '#2cc6ff'; // Blue for dragged keyframe
+                    ctx.strokeStyle = '#ffffff';
+                    const radius = 7; // Larger radius for dragged keyframe
+                    ctx.beginPath();
+                    ctx.arc(keyX, timelineMidY, radius, 0, Math.PI * 2);
+                    ctx.fill();
+                    ctx.lineWidth = 2;
                     ctx.stroke();
+                } else {
+                    ctx.fillStyle = '#ff8c00';
+                    ctx.strokeStyle = '#ffffff';
+                    const radius = currentFrame === keyFrame ? 5 : 4;
+                    ctx.beginPath();
+                    ctx.arc(keyX, timelineMidY, radius, 0, Math.PI * 2);
+                    ctx.fill();
+                    if (currentFrame === keyFrame) {
+                        ctx.lineWidth = 1.5;
+                        ctx.stroke();
+                    }
                 }
             }
             ctx.restore();
@@ -449,6 +467,26 @@ export class BoxLayerWidget extends RgthreeBaseWidget {
     }
 
     onMouseUp(event, pos, node) {
+        // Finalize keyframe drag if active
+        if (this._keyframeDragActive && this._keyframeDragStartFrame !== null) {
+            const editor = node?.editor;
+            const targetFrame = this._calculateTimelinePoint(node, pos[0]);
+
+            if (targetFrame != null && editor?.moveBoxLayerKey) {
+                // Move the keyframe (this will overwrite any existing keyframe at target)
+                editor.moveBoxLayerKey(this, this._keyframeDragStartFrame, targetFrame);
+            }
+
+            // Reset drag state
+            this._keyframeDragActive = false;
+            this._keyframeDragStartFrame = null;
+
+            if (node?.editor?.clearBoxTimelinePreview) {
+                node.editor.clearBoxTimelinePreview(this);
+            }
+            node?.setDirtyCanvas?.(true, true);
+        }
+
         if (node?.editor?.clearBoxTimelinePreview) {
             node.editor.clearBoxTimelinePreview(this);
         }
@@ -661,6 +699,24 @@ export class BoxLayerWidget extends RgthreeBaseWidget {
             }
         }
 
+        // D+Click: Start dragging a keyframe
+        if (event?.button === 0 && editor?.isShortcutActive?.('d')) {
+            event.preventDefault?.();
+            event.stopPropagation?.();
+            const targetFrame = this._calculateTimelinePoint(node, pos[0]);
+            if (targetFrame != null && this._hasKeyAtFrame(node, targetFrame)) {
+                // Start dragging this keyframe
+                this._keyframeDragActive = true;
+                this._keyframeDragStartFrame = targetFrame;
+                this._timelineDragging = false;
+                this._timelinePreviewMode = false;
+                editor?.clearBoxTimelinePreview?.(this);
+                node?.setDirtyCanvas?.(true, true);
+                return true;
+            }
+            // If D+click but not on a keyframe, fall through to normal timeline behavior
+        }
+
         if (event?.button === 2) {
             const targetFrame = this._getTimelinePoint(node);
             if (this._hasKeyAtFrame(node, targetFrame) && node?.editor?.deleteBoxLayerKey) {
@@ -683,6 +739,16 @@ export class BoxLayerWidget extends RgthreeBaseWidget {
     }
 
     onTimelineMove(event, pos, node) {
+        // Handle keyframe dragging (D+drag)
+        if (this._keyframeDragActive && this._keyframeDragStartFrame !== null) {
+            const newFrame = this._calculateTimelinePoint(node, pos[0]);
+            if (newFrame != null && newFrame !== this._keyframeDragStartFrame) {
+                // Preview: update the keyframe position temporarily
+                this._previewKeyframeMove(node, this._keyframeDragStartFrame, newFrame);
+            }
+            return true;
+        }
+
         if (!this._timelineDragging) return false;
         this._setTimelinePointFromPosition(pos[0], node, !this._timelinePreviewMode);
         return true;
@@ -990,4 +1056,24 @@ export class BoxLayerWidget extends RgthreeBaseWidget {
         }
     }
 
+    _previewKeyframeMove(node, fromFrame, toFrame) {
+        // Create a temporary preview of the keyframe at the new position
+        // This updates the timeline point and shows what the move would look like
+        const editor = node?.editor;
+        if (!editor) return;
+
+        const maxFrames = this._getMaxFrames(node);
+        const toFrameClamped = Math.max(1, Math.min(maxFrames, Math.round(toFrame)));
+
+        // Update timeline point to show preview
+        this.value.box_timeline_point = toFrameClamped;
+
+        // Use the preview system to show what the box would look like at the new frame
+        if (editor?.setBoxTimelinePreview) {
+            editor.setBoxTimelinePreview(this, toFrameClamped);
+        }
+
+        node?.setDirtyCanvas?.(true, true);
+    }
+
 }

web/power_spline_editor/prepare_refs.js

@@ -5,6 +5,7 @@ import { chainCallback, hideWidgetForGood } from './general_utils.js';
 import { getReferenceImageFromConnectedNode } from './graph_query.js';
 import { RefLayerWidget } from './layer_type_ref.js';
 import { attachLassoHelpers } from './canvas_lasso.js';
+import { attachSam2Helpers } from './canvas/canvas_sam2.js';
 import RefCanvas from './canvas_main_ref.js';
 
 // ===================================================================
@@ -18,6 +19,21 @@ function safeGetSessionItem(key) {
   try { return sessionStorage.getItem(key); } catch (e) { return null; }
 }
 
+// ===================================================================
+// Add SAM2 Magic Wand Button Styles
+// ===================================================================
+const sam2Style = document.createElement('style');
+sam2Style.textContent = `
+  .magic-wand-btn.active {
+    background: #4CAF50 !important;
+    color: white !important;
+  }
+  .magic-wand-btn:hover {
+    background: #555 !important;
+  }
+`;
+document.head.appendChild(sam2Style);
+
 // ===================================================================
 // Image Loading Utilities
 // ===================================================================
@@ -237,6 +253,7 @@ app.registerExtension({
     nodeType.prototype.setupCanvas = function () {
       this.refCanvasEditor = new RefCanvas(this.refsCanvas, this);
       attachLassoHelpers(this);
+      attachSam2Helpers(this);
 
       this.refreshCanvas = () => {
         this.refCanvasEditor.render();
@@ -261,6 +278,11 @@ app.registerExtension({
         });
 
         this.renderLassoPreview?.();
+
+        // Render SAM2 visuals if in SAM2 mode
+        if (this._sam2Mode) {
+          this.renderSam2Visuals?.(ctx, width, height);
+        }
       };
 
       this.forceCanvasRefresh = () => {
@@ -329,6 +351,23 @@ app.registerExtension({
     };
 
     nodeType.prototype.addRefLayer = function () {
+      // Exit SAM2 mode if active before creating new layer
+      if (this._sam2Mode) {
+        this._sam2Mode = false;
+        const oldLayer = this._sam2ActiveLayer;
+        this._sam2ActiveLayer = null;
+        this._sam2Points = [];
+        this._sam2MaskPath = null;
+        this.clearSam2Visuals?.();
+
+        // Update wand button state on old layer if it exists
+        if (oldLayer && oldLayer.magicWandButton) {
+          oldLayer.magicWandButton.classList.remove('active');
+          oldLayer.magicWandButton.style.background = 'transparent';
+          oldLayer.magicWandButton.style.color = '#888';
+        }
+      }
+
       const nextNum = (this.refLayers.reduce((m, l) => {
         const n = parseInt(l.value.name.match(/ref_(\d+)/)?.[1] || 0);
         return Math.max(m, n);
@@ -357,13 +396,27 @@ app.registerExtension({
         cursor: pointer; user-select: none; transition: background 0.1s ease;
       `;
 
+      // 1. Delete button (first)
+      const delBtn = this.createIconButton('Delete', '✕', () => this.removeRefLayer(layer));
+      el.appendChild(delBtn);
+
+      // 2. Name (second)
       const name = document.createElement('span');
       name.textContent = layer.value.name;
       name.style.flex = '1';
+      name.style.marginLeft = '6px'; // Add spacing after delete button
       el.appendChild(name);
 
+      // 3. Pencil/Edit button (third)
       const editBtn = this.createIconButton('Edit', '✎', () => {
         this.selectRefLayer(layer);
+
+        // If in SAM2 mode, exit SAM2 first (bakes mask to pencil)
+        if (this._sam2Mode) {
+          this.exitSam2Mode?.();
+        }
+
+        // Then handle pencil mode
         if (this._lassoDrawingActive && this._lassoActiveLayer === layer) {
           this.exitLassoMode?.();
         } else {
@@ -371,12 +424,20 @@ app.registerExtension({
         }
       });
       el.appendChild(editBtn);
-      
+
       // Store reference to edit button for color updates
       layer.editButton = editBtn;
 
-      const delBtn = this.createIconButton('Delete', '✕', () => this.removeRefLayer(layer));
-      el.appendChild(delBtn);
+      // 4. Magic Wand button (last)
+      const magicWandBtn = this.createIconButton('Magic Wand (SAM2)', '🪄', () => {
+        this.selectRefLayer(layer);
+        this.toggleSam2Mode?.();
+      });
+      magicWandBtn.classList.add('magic-wand-btn');
+      el.appendChild(magicWandBtn);
+
+      // Store reference to magic wand button for state updates
+      layer.magicWandButton = magicWandBtn;
 
       el.onclick = (e) => {
         if (e.target.tagName === 'BUTTON') return;
@@ -445,6 +506,22 @@ app.registerExtension({
       const idx = this.refLayers.indexOf(layer);
       if (idx === -1) return;
 
+      // Clear SAM2 mode if deleting the active SAM2 layer
+      if (this._sam2Mode && this._sam2ActiveLayer === layer) {
+        this._sam2Mode = false;
+        this._sam2ActiveLayer = null;
+        this._sam2Points = [];
+        this._sam2MaskPath = null;
+        this.clearSam2Visuals?.();
+
+        // Update wand button state if it exists
+        if (layer.magicWandButton) {
+          layer.magicWandButton.classList.remove('active');
+          layer.magicWandButton.style.background = 'transparent';
+          layer.magicWandButton.style.color = '#888';
+        }
+      }
+
       if (this.selectedRefLayer === layer) this.selectedRefLayer = null;
       this.refLayers.splice(idx, 1);
       if (layer.domElement?.parentNode) layer.domElement.parentNode.removeChild(layer.domElement);