initial support for matanyone2

README.md

@@ -4,9 +4,17 @@ MatAnyone in ComfyUI (Remove background)
 
 Stable Video Matting with Consistent Memory Propagation: <https://github.com/pq-yang/MatAnyone>
 
+Scaling Video Matting via a Learned Quality Evaluator: <https://github.com/pq-yang/MatAnyone2>
+
 Download `matanyone.pth` from <https://github.com/pq-yang/MatAnyone?tab=readme-ov-file#download-model>
 
-and place it in `checkpoint/matanyone.pth`
+Download `matanyone2.pth` from <https://github.com/pq-yang/MatAnyone2?tab=readme-ov-file#-inference>
+
+```bash
+checkpoint/
+    matanyone.pth
+    matanyone2.pth
+```
 
 ## Workflow
 
@@ -25,6 +33,19 @@ Inputs:
 
 Your input mask won't actually be in the final matte. Instead, the warmup process generate a new input mask, which is then propagated throughout the video.
 
+### MatAnyone 2 Support
+
+The extension now includes the `MatAnyone2Video` node, which runs the improved MatAnyoneV2 model for higher-quality and more robust video matting.
+
+**Workflow V2**: [workflow/workflow_mat_anyone_v2.json](workflow/workflow_mat_anyone_v2.json)
+
+![matanyone_example](workflow/mat_anyone_preview_v2.png)
+(Not a workflow-embedded image)
+
+Additional Inputs for V2:
+- `r_erode` (optional): The radius for morphological erosion applied to the `foreground_mask` before processing (defaults to 0). Useful for refining rough masks.
+- `r_dilate` (optional): The radius for morphological dilation applied to the `foreground_mask` before processing (defaults to 0).
+
 ## Credit
 
 ```cite
@@ -34,4 +55,11 @@ Your input mask won't actually be in the final matte. Instead, the warmup proces
     booktitle = {arXiv preprint arXiv:2501.14677},
     year      = {2025}
 }
+
+@InProceedings{yang2026matanyone2,
+   title     = {{MatAnyone 2}: Scaling Video Matting via a Learned Quality Evaluator},
+   author    = {Yang, Peiqing and Zhou, Shangchen and Hao, Kai and Tao, Qingyi},
+   booktitle = {CVPR},
+   year      = {2026}
+   }
 ```

__init__.py

@@ -1,12 +1,14 @@
-from .run import MatAnyoneVideo, SolidColorBatched
+from .run import MatAnyoneVideo, MatAnyone2Video, SolidColorBatched
 
 NODE_CLASS_MAPPINGS = {
     "MatAnyone": MatAnyoneVideo,
+    "MatAnyone2": MatAnyone2Video,
     "SolidColorBatched": SolidColorBatched,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
     "MatAnyone": "MatAnyone",
+    "MatAnyone2": "MatAnyone2",
     "SolidColorBatched": "Solid Color Batched",
 }
 

constants.py

@@ -1 +1,2 @@
 ckpt_path = "checkpoint/matanyone.pth"
+ckpt_path_matanyone2 = "checkpoint/matanyone2.pth"

run.py

@@ -18,7 +18,12 @@ from .mat_anyone import (
     get_matanyone_model,
     inference_matanyone,
 )
+from .mat_anyone2 import (
+    get_matanyone2_model_cached,
+    inference_matanyone2,
+)
 from .src.core.inference_core import InferenceCore
+from .src.matanyone2.inference.inference_core import InferenceCore as InferenceCore2
 
 
 class MatAnyoneVideo:
@@ -91,6 +96,91 @@ class MatAnyoneVideo:
         )
 
 
+class MatAnyone2Video:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "src_video": ("IMAGE",),
+                "mask_frame": (
+                    "INT",
+                    {"default": 0, "min": 0, "step": 1},
+                ),
+                "n_warmup": (
+                    "INT",
+                    {"default": 10, "min": 1, "step": 1},
+                ),
+            },
+            "optional": {
+                "foreground_mask": ("IMAGE",),
+                "foreground_MASK": ("MASK",),
+                "solid_color": ("IMAGE",),
+                "r_erode": (
+                    "INT",
+                    {"default": 0, "min": 0, "step": 1},
+                ),
+                "r_dilate": (
+                    "INT",
+                    {"default": 0, "min": 0, "step": 1},
+                ),
+            },
+        }
+
+    RETURN_TYPES = (
+        "IMAGE",
+        "IMAGE",
+    )
+    RETURN_NAMES = (
+        "matte",
+        "green_screen",
+    )
+    FUNCTION = "process"
+    CATEGORY = "MatAnyone"
+
+    def process(
+        self,
+        src_video: torch.Tensor,
+        mask_frame: int,
+        n_warmup: int,
+        foreground_mask: torch.Tensor | None = None,
+        foreground_MASK: torch.Tensor | None = None,
+        solid_color: torch.Tensor | None = None,
+        r_erode: int = 0,
+        r_dilate: int = 0,
+    ):
+        mask = get_mask(foreground_mask, foreground_MASK)
+        src_video = src_video.permute(0, 3, 1, 2)  # T C H W RGB, [0, 1]
+
+        matanyone2 = get_matanyone2_model_cached()
+        processor = InferenceCore2(matanyone2, cfg=matanyone2.cfg)
+        phas = inference_matanyone2(
+            src_video,
+            mask,
+            processor,
+            mask_frame=mask_frame,
+            n_warmup=n_warmup,
+            r_erode=r_erode,
+            r_dilate=r_dilate,
+        )
+        out_mask = torch.cat(phas).permute(0, 2, 3, 1)  # T H W 1 -> T H W C
+        out_mask_rgb = out_mask.repeat(1, 1, 1, 3)
+
+        gb = torch.empty(
+            0, out_mask_rgb.shape[1], out_mask_rgb.shape[2], out_mask_rgb.shape[3]
+        )
+        if solid_color is not None:
+            src_video_hwc = src_video.permute(0, 2, 3, 1).to(out_mask_rgb.device)
+            solid_color_batched = solid_color.repeat(
+                src_video_hwc.shape[0], 1, 1, 1
+            ).to(out_mask_rgb.device)
+            gb = out_mask_rgb * src_video_hwc + (1 - out_mask_rgb) * solid_color_batched
+
+        return (
+            out_mask_rgb,
+            gb,
+        )
+
+
 class SolidColorBatched:
     @classmethod
     def INPUT_TYPES(s):

src/matanyone2/config/__init__.py

@@ -0,0 +1 @@
+"""MatAnyone2 config."""

src/matanyone2/config/eval_matanyone_config.yaml

@@ -0,0 +1,33 @@
+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/matanyone2.pth
+output_dir: null
+flip_aug: False
+max_internal_size: -1
+save_all: True
+use_all_masks: False
+use_long_term: False
+mem_every: 5
+max_mem_frames: 5
+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
+save_scores: False
+save_aux: False
+visualize: False

src/matanyone2/config/model/base.yaml

@@ -0,0 +1,49 @@
+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]
+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:
+  up_dims: [256, 128, 128, 64, 16]

src/matanyone2/inference/__init__.py

@@ -0,0 +1 @@
+"""MatAnyone2 inference modules."""

src/matanyone2/inference/image_feature_store.py

@@ -0,0 +1,59 @@
+"""
+Image feature cache for MatAnyone2 inference.
+"""
+
+import warnings
+
+import torch
+
+from ..model.matanyone2 import MatAnyone2
+
+
+class ImageFeatureStore:
+    """Caches image features for reuse during inference."""
+
+    def __init__(self, network: MatAnyone2, 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) -> tuple:
+        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) -> tuple:
+        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) -> tuple:
+        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(
+                "Leaking %s in the image feature store" % list(self._store.keys())
+            )

src/matanyone2/inference/inference_core.py

@@ -0,0 +1,381 @@
+"""
+Inference core for MatAnyone2.
+"""
+
+import logging
+from typing import Iterable, List, Optional, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from omegaconf import DictConfig
+
+from ..model.matanyone2 import MatAnyone2
+from ..utils.device import get_default_device, safe_autocast_decorator
+from ..utils.tensor_utils import aggregate, pad_divide_by, unpad
+from .image_feature_store import ImageFeatureStore
+from .memory_manager import MemoryManager
+from .object_manager import ObjectManager
+
+log = logging.getLogger()
+
+
+class InferenceCore:
+    """Core inference logic for MatAnyone2."""
+
+    def __init__(
+        self,
+        network: MatAnyone2,
+        cfg: Optional[DictConfig] = None,
+        *,
+        image_feature_store: ImageFeatureStore = None,
+        device: Optional[Union[str, torch.device]] = None,
+    ):
+        if device is None:
+            device = get_default_device()
+        self.device = device
+        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
+        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) -> None:
+        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) -> None:
+        self.curr_ti = -1
+        self.last_mem_ti = 0
+        self.memory.clear_non_permanent_memory()
+
+    def clear_sensory_memory(self) -> None:
+        self.curr_ti = -1
+        self.last_mem_ti = 0
+        self.memory.clear_sensory_memory()
+
+    def update_config(self, cfg) -> None:
+        self.mem_every = cfg["mem_every"]
+        self.memory.update_config(cfg)
+
+    def clear_temp_mem(self) -> None:
+        self.memory.clear_work_mem()
+        self.memory.clear_obj_mem()
+
+    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:
+        if prob.shape[1] == 0:
+            log.warning("Trying to add an empty object mask to memory!")
+            return
+
+        as_permanent = "all" if force_permanent else "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:
+        bs = key.shape[0]
+        if not self.memory.engaged:
+            log.warning("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:
+            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,
+        )
+        if self.flip_aug:
+            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
+
+    @torch.inference_mode()
+    @safe_autocast_decorator()
+    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:
+        if objects is None and mask is not None:
+            assert not idx_mask
+            objects = list(range(1, mask.shape[0] + 1))
+
+        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)
+        image = image.unsqueeze(0)
+        if self.flip_aug:
+            image = torch.cat([image, torch.flip(image, dims=[-1])], dim=0)
+
+        is_mem_frame = (
+            (self.curr_ti - self.last_mem_ti >= self.mem_every) or mask is not None
+        ) and (not end)
+        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
+        )
+
+        if first_frame_pred:
+            self.curr_ti = 0
+            self.last_mem_ti = 0
+            is_mem_frame = True
+            need_segment = True
+            update_sensory = True
+
+        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
+        )
+
+        if need_segment:
+            pred_prob_with_bg = self._segment(
+                key, selection, pix_feat, ms_feat, update_sensory=update_sensory
+            )
+
+        if mask is not None:
+            corresponding_tmp_ids, _ = self.object_manager.add_new_objects(objects)
+            mask, _ = pad_divide_by(mask, 16)
+            if need_segment:
+                pred_prob_no_bg = pred_prob_with_bg[1:]
+                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:
+                        pred_prob_no_bg[tmp_id - 1] = this_mask
+                mask = torch.cat([pred_prob_no_bg, *new_masks], dim=0)
+            elif idx_mask:
+                if len(objects) == 0:
+                    if delete_buffer:
+                        self.image_feature_store.delete(self.curr_ti)
+                    log.warning("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.0
+                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)
+
+        self.last_mask = pred_prob_with_bg[1:].unsqueeze(0)
+        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
+
+        if is_mem_frame or force_permanent:
+            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:
+            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:
+            output_prob = F.interpolate(
+                output_prob.unsqueeze(0),
+                size=(h, w),
+                mode="bilinear",
+                align_corners=False,
+            )[0]
+
+        return output_prob
+
+    def output_prob_to_mask(
+        self, output_prob: torch.Tensor, matting: bool = True
+    ) -> torch.Tensor:
+        if matting:
+            return output_prob[1:].squeeze(0)
+        mask = torch.argmax(output_prob, dim=0)
+        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

src/matanyone2/inference/kv_memory_store.py

@@ -0,0 +1,326 @@
+"""
+Key-value memory store for MatAnyone2 inference.
+"""
+
+from collections import defaultdict
+from typing import Dict, List, Literal, Optional
+
+import torch
+
+
+def _add_last_dim(dictionary, key, new_value, prepend=False):
+    """Appends or prepends a new value to the last dimension of a tensor."""
+    if key in dictionary:
+        dictionary[key] = torch.cat([dictionary[key], new_value], -1)
+    else:
+        dictionary[key] = new_value
+
+
+class KeyValueMemoryStore:
+    """Key-value memory store for working and long-term memory."""
+
+    def __init__(self, save_selection: bool = False, save_usage: bool = False):
+        self.save_selection = save_selection
+        self.save_usage = save_usage
+        self.global_bucket_id = 0
+        self.buckets: Dict[int, List[int]] = {}
+        self.k: Dict[int, torch.Tensor] = {}
+        self.v: Dict[int, torch.Tensor] = {}
+        self.perm_end_pt: Dict[int, int] = defaultdict(int)
+        self.s = {}
+        if self.save_selection:
+            self.e = {}
+        if self.save_usage:
+            self.use_cnt = {}
+            self.life_cnt = {}
+
+    @property
+    def key(self):
+        return self.k
+
+    @property
+    def value(self):
+        return self.v
+
+    @property
+    def shrinkage(self):
+        return self.s
+
+    @property
+    def selection(self):
+        return self.e if self.save_selection else None
+
+    def add(
+        self,
+        key: torch.Tensor,
+        values: Dict[int, torch.Tensor],
+        shrinkage: torch.Tensor,
+        selection: Optional[torch.Tensor] = None,
+        supposed_bucket_id: int = -1,
+        as_permanent: Literal["no", "first", "all"] = "no",
+    ) -> None:
+        bs = key.shape[0]
+        ne = key.shape[-1]
+        assert len(key.shape) == 3
+        assert len(shrinkage.shape) == 3
+        assert not self.save_selection or (
+            selection is not None and len(selection.shape) == 3
+        )
+        assert as_permanent in ["no", "first", "all"]
+
+        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 = [
+                        bid for bid, oids in self.buckets.items() if obj in oids
+                    ]
+                    assert len(bucket_used) == 1
+                    enabled_buckets.add(bucket_used[0])
+                else:
+                    self.v[obj] = value
+                    if new_bucket_id is None:
+                        new_bucket_id = self.global_bucket_id
+                        self.global_bucket_id += 1
+                        self.buckets[new_bucket_id] = []
+                    self.buckets[new_bucket_id].append(obj)
+                    enabled_buckets.add(new_bucket_id)
+
+        add_as_permanent = {}
+        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
+
+        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
+            )
+
+        for bucket_id in self.buckets:
+            if bucket_id not in enabled_buckets:
+                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 and selection is not None:
+                    _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:
+        if not self.save_usage:
+            return
+        usage = usage[:, self.perm_end_pt[bucket_id] :]
+        if usage.shape[-1] == 0:
+            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:
+        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:
+            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]
+        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:
+            e = self.e[bucket_id]
+            self.e[bucket_id] = torch.cat(
+                [e[:, :, : start - p_size], e[:, :, end:]], -1
+            )
+        if self.save_usage:
+            use_cnt = self.use_cnt[bucket_id]
+            life_cnt = self.life_cnt[bucket_id]
+            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:
+        object_ids = self.buckets[bucket_id]
+        assert self.perm_end_pt[bucket_id] == 0
+        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())
+        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:
+            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:
+        if not self.save_usage:
+            raise RuntimeError("Usage not tracked")
+        return self.use_cnt[bucket_id] / self.life_cnt[bucket_id]
+
+    def get_all_sliced(self, bucket_id: int, start: int, end: int) -> tuple:
+        assert start >= 0
+        assert end <= 0
+        p_size = self.perm_end_pt[bucket_id]
+        start = start + p_size
+        if end == 0:
+            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]) -> None:
+        obj_keep_idx = set(obj_keep_idx)
+        buckets_to_remove = []
+        for bucket_id, object_ids in self.buckets.items():
+            self.buckets[bucket_id] = [oid for oid in object_ids if oid in obj_keep_idx]
+            if len(self.buckets[bucket_id]) == 0:
+                buckets_to_remove.append(bucket_id)
+        self.v = {k: v for k, v in self.v.items() if k in obj_keep_idx}
+        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) -> None:
+        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
+        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
+        return bucket_id in self.buckets
+
+    @property
+    def num_objects(self) -> int:
+        return len(self.v)

src/matanyone2/inference/memory_manager.py

@@ -0,0 +1,421 @@
+"""
+Memory manager for MatAnyone2 inference.
+"""
+
+import logging
+from typing import Dict, List
+
+import torch
+from omegaconf import DictConfig
+
+from ..model.matanyone2 import MatAnyone2
+from ..model.utils.memory_utils import do_softmax, get_similarity
+from .kv_memory_store import KeyValueMemoryStore
+from .object_manager import ObjectManager
+
+log = logging.getLogger()
+
+
+class MemoryManager:
+    """Manages working and long-term memory for MatAnyone2."""
+
+    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
+
+        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
+
+        self.CK = self.CV = None
+        self.H = self.W = None
+        self.sensory = {}
+        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
+        assert 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:
+        if len(v.shape) == 3:
+            if uncert_mask is not None:
+                return v @ affinity * uncert_mask
+            return v @ affinity
+        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:
+        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:
+        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: torch.Tensor,
+        pix_feat: torch.Tensor,
+        last_mask: torch.Tensor,
+        network: MatAnyone2,
+        uncert_output=None,
+    ) -> Dict[int, torch.Tensor]:
+        bs = pix_feat.shape[0]
+        all_readout_mem = {}
+        buckets = self.work_mem.buckets
+
+        for bucket_id, bucket in buckets.items():
+            object_chunks = (
+                [bucket]
+                if self.chunk_size < 1
+                else [
+                    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)
+                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:
+                    self.aux = {
+                        "q_logits": (aux_features["logits"] if aux_features else None),
+                    }
+        return all_readout_mem
+
+    def read(
+        self,
+        pix_feat: torch.Tensor,
+        query_key: torch.Tensor,
+        selection: torch.Tensor,
+        last_mask: torch.Tensor,
+        network: MatAnyone2,
+        uncert_output=None,
+        last_msk_value=None,
+        ti=None,
+        last_pix_feat=None,
+        last_pred_mask=None,
+    ) -> Dict[int, torch.Tensor]:
+        h, w = pix_feat.shape[-2:]
+        bs = pix_feat.shape[0]
+        uncert_mask = uncert_output["mask"] if uncert_output is not None else None
+        query_key = query_key.flatten(start_dim=2)
+        selection = selection.flatten(start_dim=2)
+
+        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):
+                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,
+                )
+                work_usage = usage[:, long_mem_size:]
+                self.work_mem.update_bucket_usage(bucket_id, work_usage)
+                if self.count_long_term_usage:
+                    long_usage = usage[:, :long_mem_size]
+                    self.long_mem.update_bucket_usage(bucket_id, long_usage)
+            else:
+                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)
+
+            object_chunks = (
+                [bucket]
+                if self.chunk_size < 1
+                else [
+                    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)
+                visual_readout = self._readout(
+                    affinity, this_msk_value, uncert_mask
+                ).view(bs, len(objects), self.CV, h, w)
+
+                if (
+                    last_msk_value is not None
+                    and last_pix_feat is not None
+                    and last_pred_mask is not None
+                ):
+                    uncert_output = network.pred_uncertainty(
+                        last_pix_feat,
+                        pix_feat,
+                        last_pred_mask,
+                        visual_readout[:, 0] - last_msk_value[:, 0],
+                    )
+                    uncert_prob = uncert_output["prob"].unsqueeze(1)
+                    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:
+                    self.aux = {
+                        "q_logits": (aux_features["logits"] if aux_features else None),
+                    }
+        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:
+        bs = key.shape[0]
+        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
+            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 = 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:
+            selection = selection.flatten(start_dim=2)
+
+        for obj_id, obj in enumerate(objects):
+            if obj in self.obj_v:
+                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]
+
+        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 list(self.work_mem.buckets.keys()):
+            if self.use_long_term:
+                if self.work_mem.non_perm_size(bucket_id) >= self.max_work_tokens:
+                    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:
+                self.work_mem.remove_old_memory(bucket_id, self.max_work_tokens)
+
+    def purge_except(self, obj_keep_idx: List[int]) -> None:
+        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():
+            self.engaged = False
+
+    def compress_features(self, bucket_id: int) -> None:
+        k, sk, ek, value, usage = self.work_mem.get_all_sliced(
+            bucket_id, 0, -self.min_work_tokens
+        )
+        prototype_key, prototype_value, prototype_shrinkage = self.consolidation(
+            k, sk, ek, value, usage
+        )
+        self.work_mem.sieve_by_range(
+            bucket_id, 0, -self.min_work_tokens, min_size=self.min_work_tokens
+        )
+        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,
+    ) -> tuple:
+        bs = candidate_key.shape[0]
+        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])
+            if candidate_selection is not None:
+                prototype_selection.append(
+                    candidate_selection[bi, :, prototype_indices]
+                )
+        prototype_key = torch.stack(prototype_key, dim=0)
+        prototype_selection = (
+            torch.stack(prototype_selection, dim=0) if prototype_selection else None
+        )
+        similarity = get_similarity(
+            candidate_key,
+            candidate_shrinkage,
+            prototype_key,
+            prototype_selection,
+        )
+        affinity = do_softmax(similarity)
+        prototype_value = {
+            k: self._readout(affinity, v) for k, v in candidate_value.items()
+        }
+        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]
+    ) -> None:
+        for obj in ids:
+            if obj not in self.sensory:
+                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]) -> None:
+        for obj_id, obj in enumerate(ids):
+            self.sensory[obj] = sensory[:, obj_id]
+
+    def get_sensory(self, ids: List[int]):
+        return self._get_sensory_by_ids(ids)
+
+    def clear_non_permanent_memory(self) -> None:
+        self.work_mem.clear_non_permanent_memory()
+        if self.use_long_term:
+            self.long_mem.clear_non_permanent_memory()
+
+    def clear_sensory_memory(self) -> None:
+        self.sensory = {}
+
+    def clear_work_mem(self) -> None:
+        self.work_mem = KeyValueMemoryStore(
+            save_selection=self.use_long_term,
+            save_usage=self.use_long_term,
+        )
+
+    def clear_obj_mem(self) -> None:
+        self.obj_v = {}

src/matanyone2/inference/object_info.py

@@ -0,0 +1,28 @@
+"""
+Object metadata for MatAnyone2 inference.
+"""
+
+
+class ObjectInfo:
+    """Stores meta information for an object."""
+
+    def __init__(self, id: int):
+        self.id = id
+        self.poke_count = 0
+
+    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 "(ID: %s)" % self.id

src/matanyone2/inference/object_manager.py

@@ -0,0 +1,134 @@
+"""
+Object manager for MatAnyone2 inference.
+"""
+
+from typing import Dict, List, Union
+
+import torch
+
+from .object_info import ObjectInfo
+
+
+class ObjectManager:
+    """
+    Manages object IDs and temporary IDs.
+    Object IDs are immutable; temporary IDs indicate tensor positions.
+    """
+
+    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]],
+    ) -> tuple:
+        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:
+                corresponding_tmp_ids.append(self.obj_to_tmp_id[obj])
+                corresponding_obj_ids.append(obj.id)
+            else:
+                new_obj = ObjectInfo(id=obj.id)
+                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(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:
+        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) -> tuple:
+        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:
+        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]:
+        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:
+        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])
+        return torch.stack(output, dim=dim)
+
+    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_id_to_obj:
+                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]]

src/matanyone2/model/aux_modules.py

@@ -0,0 +1,98 @@
+"""
+Auxiliary output modules for MatAnyone2.
+"""
+
+from typing import Dict
+
+import torch
+import torch.nn as nn
+from omegaconf import DictConfig
+
+from ..utils.tensor_utils import aggregate
+from .group_modules import GConv2d
+
+
+class LinearPredictor(nn.Module):
+    """Linear predictor for auxiliary loss."""
+
+    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:
+        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):
+    """Direct predictor for auxiliary loss."""
+
+    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:
+        return self.projection(x).squeeze(2)
+
+
+class AuxComputer(nn.Module):
+    """Computes auxiliary outputs for auxiliary losses."""
+
+    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
+        return aggregate(prob, dim=1)
+
+    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:
+            logits = self.sensory_aux(pix_feat, sensory)
+            aux_output["sensory_logits"] = self._aggregate_with_selector(
+                logits, selector
+            )
+        if self.use_query_aux:
+            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]:
+        q_logits = aux_input["q_logits"]
+        aux_output = {}
+        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

src/matanyone2/model/big_modules.py

@@ -0,0 +1,402 @@
+"""
+High-level network blocks for MatAnyone2.
+"""
+
+from typing import Iterable
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from omegaconf import DictConfig
+
+from ..utils.device import safe_autocast
+from .group_modules import GConv2d, GroupFeatureFusionBlock, MainToGroupDistributor
+from .modules import (
+    DecoderFeatureProcessor,
+    MaskUpsampleBlock,
+    SensoryDeepUpdater,
+    SensoryUpdater_fullscale,
+)
+from .utils import resnet
+
+
+class UncertPred(nn.Module):
+    """Uncertainty prediction 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
+
+    def train(self, mode=True):
+        self.training = False
+        for module in self.children():
+            module.train(False)
+        return self
+
+
+class PixelEncoder(nn.Module):
+    """Pixel/feature encoder (ResNet backbone)."""
+
+    def __init__(self, model_cfg: DictConfig):
+        super().__init__()
+        self.is_resnet = "resnet" in model_cfg.pixel_encoder.type
+        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) -> tuple:
+        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
+
+    def train(self, mode=True):
+        self.training = False
+        for module in self.children():
+            module.train(False)
+        return self
+
+
+class KeyProjection(nn.Module):
+    """Key projection for memory."""
+
+    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)
+        self.d_proj = nn.Conv2d(mid_dim, 1, kernel_size=3, padding=1)
+        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) -> tuple:
+        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):
+    """Mask encoder with ResNet backbone."""
+
+    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
+
+        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,
+    ) -> tuple:
+        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:
+            new_sensory = torch.empty_like(sensory) if deep_update else sensory
+            fast_path = False
+
+        all_g = []
+        for i in range(0, num_objects, chunk_size):
+            g_chunk = g if fast_path else 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)
+            g_chunk = self.maxpool(g_chunk)
+            g_chunk = self.relu(g_chunk)
+            g_chunk = self.layer1(g_chunk)
+            g_chunk = self.layer2(g_chunk)
+            g_chunk = self.layer3(g_chunk)
+
+            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
+
+    def train(self, mode=True):
+        self.training = False
+        for module in self.children():
+            module.train(False)
+        return self
+
+
+class PixelFeatureFuser(nn.Module):
+    """Fuses pixel and sensory memory features."""
+
+    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
+
+        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):
+    """Mask decoder with multi-scale upsample."""
+
+    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])
+        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,
+    ) -> tuple:
+        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:
+            new_sensory = torch.empty_like(sensory) if update_sensory else sensory
+            fast_path = False
+
+        all_logits = []
+        for i in range(0, num_objects, chunk_size):
+            p16 = memory_readout if fast_path else 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 safe_autocast(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
+                        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
+                        logits = last_mask + res
+                    else:
+                        logits = self.pred_mat(
+                            F.relu(p1.flatten(start_dim=0, end_dim=1).float())
+                        )
+
+            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

src/matanyone2/model/channel_attn.py

@@ -0,0 +1,46 @@
+"""
+Channel attention block for MatAnyone2.
+"""
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class CAResBlock(nn.Module):
+    """Channel attention residual block."""
+
+    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()
+
+        if self.residual:
+            x = x * w + self.downsample(r)
+        else:
+            x = x * w
+
+        return x

src/matanyone2/model/group_modules.py

@@ -0,0 +1,139 @@
+"""
+Group convolution and feature fusion modules for MatAnyone2.
+"""
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .channel_attn import CAResBlock
+
+
+def interpolate_groups(
+    g: torch.Tensor, ratio: float, mode: str, align_corners: bool
+) -> torch.Tensor:
+    """Interpolates group tensor by ratio."""
+    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:
+    """Upsamples group 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:
+    """Downsamples group tensor."""
+    return interpolate_groups(g, ratio, mode, align_corners)
+
+
+class GConv2d(nn.Conv2d):
+    """Group convolution (batch, num_objects, C, H, W)."""
+
+    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):
+    """Residual block for group tensors."""
+
+    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):
+    """Distributes main features to group dimension."""
+
+    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):
+    """Fuses main and group features."""
+
+    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

src/matanyone2/model/matanyone2.py

@@ -0,0 +1,364 @@
+"""
+MatAnyone2 model - video matting network.
+"""
+
+import logging
+from typing import Dict, Iterable, List, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from omegaconf import DictConfig, OmegaConf
+
+from ..utils.device import get_default_device, safe_autocast
+from ..utils.tensor_utils import aggregate
+from .aux_modules import AuxComputer
+from .big_modules import (
+    KeyProjection,
+    MaskDecoder,
+    MaskEncoder,
+    PixelEncoder,
+    PixelFeatureFuser,
+    UncertPred,
+)
+from .transformer.object_summarizer import ObjectSummarizer
+from .transformer.object_transformer import QueryTransformer
+from .utils.memory_utils import get_affinity, readout
+
+log = logging.getLogger()
+
+
+class MatAnyone2(nn.Module):
+    """MatAnyone2 video matting model."""
+
+    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
+
+        log.info("Single object: %s", self.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:
+        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
+        return {"logits": logits, "prob": prob, "mask": mask}
+
+    def encode_image(
+        self, image: torch.Tensor, seq_length=None, last_feats=None
+    ) -> tuple:
+        device = get_default_device()
+        self.pixel_mean = self.pixel_mean.to(device)
+        self.pixel_std = self.pixel_std.to(device)
+        image = (image - self.pixel_mean) / self.pixel_std
+        ms_image_feat = self.pixel_encoder(image, seq_length)
+        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
+
+    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], dict]:
+        batch_size, num_objects = msk_value.shape[:2]
+        uncert_mask = uncert_output["mask"] if uncert_output is not None else None
+
+        with safe_autocast(enabled=False):
+            affinity = get_affinity(
+                memory_key.float(),
+                memory_shrinkage.float(),
+                query_key.float(),
+                query_selection.float(),
+                uncert_mask=uncert_mask,
+            )
+            msk_value_flat = msk_value.flatten(start_dim=1, end_dim=2).float()
+            pixel_readout = readout(affinity, msk_value_flat, uncert_mask)
+            pixel_readout = pixel_readout.view(
+                batch_size, num_objects, self.value_dim, *pixel_readout.shape[-2:]
+            )
+
+            last_val = msk_value[:, :, :, -1]
+            mem_val_diff = pixel_readout[:, 0] - last_val[:, 0]
+            uncert_output = self.pred_uncertainty(
+                last_pix_feat, pix_feat, last_pred_mask, mem_val_diff
+            )
+            uncert_prob = uncert_output["prob"].unsqueeze(1)
+            pixel_readout = pixel_readout * uncert_prob + last_val.unsqueeze(2) * (
+                1 - uncert_prob
+            )
+
+        pixel_readout = self.pixel_fusion(pix_feat, pixel_readout, sensory, last_mask)
+        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: torch.Tensor,
+        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]]:
+        pixel_readout = self.pixel_fusion(pix_feat, pixel_readout, sensory, last_mask)
+        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)
+        return self.pixel_fuser(
+            pix_feat,
+            pixel,
+            sensory,
+            last_mask,
+            last_others,
+            chunk_size=chunk_size,
+        )
+
+    def readout_query(
+        self,
+        pixel_readout: torch.Tensor,
+        obj_memory: torch.Tensor,
+        *,
+        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=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]:
+        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
+            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:
+            for k in list(src_dict.keys()):
+                if k == "mask_encoder.conv1.weight":
+                    if src_dict[k].shape[1] == 4:
+                        log.info(
+                            "Converting %s from single object to multiple objects.",
+                            k,
+                        )
+                        pads = torch.zeros((64, 1, 7, 7), device=src_dict[k].device)
+                        if not init_as_zero_if_needed:
+                            nn.init.orthogonal_(pads)
+                        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(
+                            "Converting %s from single object to multiple objects.",
+                            k,
+                        )
+                        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)
+                        src_dict[k] = torch.cat([src_dict[k], pads], 1)
+        elif self.single_object:
+            if src_dict["mask_encoder.conv1.weight"].shape[1] == 5:
+                log.warning(
+                    "Converting mask_encoder.conv1.weight from multiple "
+                    "objects to single object."
+                )
+                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("Key %s found in src_dict but not in self.state_dict()", k)
+        for k in self.state_dict():
+            if k not in src_dict:
+                log.info("Key %s found in self.state_dict() but not in src_dict", k)
+
+        self.load_state_dict(src_dict, strict=False)
+
+    @property
+    def device(self) -> torch.device:
+        return self.pixel_mean.device

src/matanyone2/model/modules.py

@@ -0,0 +1,173 @@
+"""
+Decoder and sensory updater modules for MatAnyone2.
+"""
+
+from typing import Iterable, List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ..utils.device import safe_autocast
+from .group_modules import (
+    GConv2d,
+    GroupResBlock,
+    MainToGroupDistributor,
+    downsample_groups,
+    upsample_groups,
+)
+
+
+class UpsampleBlock(nn.Module):
+    """Upsample block with residual connection."""
+
+    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):
+    """Mask upsample block with group features."""
+
+    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):
+    """Processes multi-scale decoder features."""
+
+    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]:
+        return [func(x) for x, func in zip(multi_scale_features, self.transforms)]
+
+
+def _recurrent_update(h: torch.Tensor, values: torch.Tensor) -> torch.Tensor:
+    """GRU-like recurrent update."""
+    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):
+    """Full-scale sensory updater for decoder."""
+
+    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 safe_autocast(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):
+    """Sensory updater for decoder (reduced scale)."""
+
+    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 safe_autocast(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):
+    """Deep sensory updater."""
+
+    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 safe_autocast(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):
+    """Residual block."""
+
+    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

src/matanyone2/model/transformer/__init__.py

@@ -0,0 +1 @@
+"""MatAnyone2 transformer modules."""

src/matanyone2/model/transformer/object_summarizer.py

@@ -0,0 +1,92 @@
+"""
+Object summarizer for MatAnyone2.
+"""
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from omegaconf import DictConfig
+
+from ...utils.device import safe_autocast
+from .positional_encoding import PositionalEncoding
+
+
+def _weighted_pooling(
+    masks: torch.Tensor, value: torch.Tensor, logits: torch.Tensor
+) -> tuple:
+    """Weighted pooling of value by logits and masks."""
+    weights = logits.sigmoid() * masks
+    sums = torch.einsum("bkhwq,bkhwc->bkqc", weights, value)
+    area = weights.flatten(start_dim=2, end_dim=3).sum(2).unsqueeze(-1)
+    return sums, area
+
+
+class ObjectSummarizer(nn.Module):
+    """Summarizes object features into fixed-size representations."""
+
+    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,
+    ) -> tuple:
+        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 safe_autocast(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
+        return summaries, None

src/matanyone2/model/transformer/object_transformer.py

@@ -0,0 +1,218 @@
+"""
+Object query transformer for MatAnyone2.
+"""
+
+from typing import Dict, Optional
+
+import torch
+import torch.nn as nn
+from omegaconf import DictConfig
+
+from ...utils.tensor_utils import aggregate
+from ..group_modules import GConv2d
+from .positional_encoding import PositionalEncoding
+from .transformer_layers import FFN, CrossAttention, PixelFFN, SelfAttention
+
+
+class QueryTransformerBlock(nn.Module):
+    """Single transformer block for object queries."""
+
+    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,
+    ) -> tuple:
+        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):
+    """Query transformer for object-level segmentation."""
+
+    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
+
+        self.query_init = nn.Embedding(self.num_queries, self.embed_dim)
+        self.query_emb = nn.Embedding(self.num_queries, self.embed_dim)
+        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,
+        )
+
+        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,
+    ) -> tuple:
+        T = obj_summaries.shape[2]
+        bs, num_objects, _, H, W = pixel.shape
+
+        obj_summaries = obj_summaries.view(
+            bs * num_objects, T, self.num_queries, self.embed_dim + 1
+        )
+        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)
+
+        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
+        )
+
+        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
+
+        aux_features = {"logits": []}
+        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
+        aux_features["p_weights"] = p_weights
+
+        if self.training:
+            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:
+        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

src/matanyone2/model/transformer/positional_encoding.py

@@ -0,0 +1,93 @@
+"""
+Positional encoding for MatAnyone2 transformer.
+"""
+
+import math
+
+import numpy as np
+import torch
+from torch import nn
+
+from ...utils.device import get_default_device
+
+
+def get_emb(sin_inp: torch.Tensor) -> torch.Tensor:
+    """Gets 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):
+    """Sinusoidal positional encoding for 2D features."""
+
+    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
+
+    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
+        if len(tensor.shape) not in (4, 5):
+            raise RuntimeError(
+                "The input tensor has to be 4/5d, got %s!" % tensor.shape
+            )
+
+        if len(tensor.shape) == 5:
+            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
+            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:
+            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
+        return self.cached_penc.unsqueeze(1)

src/matanyone2/model/transformer/transformer_layers.py

@@ -0,0 +1,182 @@
+"""
+Transformer layers for MatAnyone2.
+"""
+
+from typing import Callable, List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+from ..channel_attn import CAResBlock
+
+
+class SelfAttention(nn.Module):
+    """Self-attention with optional positional encoding."""
+
+    def __init__(
+        self,
+        dim: int,
+        nhead: int,
+        dropout: float = 0.0,
+        batch_first: bool = True,
+        add_pe_to_qkv: List[bool] = None,
+    ):
+        super().__init__()
+        if add_pe_to_qkv is None:
+            add_pe_to_qkv = [True, True, False]
+        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)
+
+
+class CrossAttention(nn.Module):
+    """Cross-attention with optional positional encoding."""
+
+    def __init__(
+        self,
+        dim: int,
+        nhead: int,
+        dropout: float = 0.0,
+        batch_first: bool = True,
+        add_pe_to_qkv: List[bool] = None,
+        residual: bool = True,
+        norm: bool = True,
+    ):
+        super().__init__()
+        if add_pe_to_qkv is None:
+            add_pe_to_qkv = [True, True, False]
+        self.cross_attn = nn.MultiheadAttention(
+            dim, nhead, dropout=dropout, batch_first=batch_first
+        )
+        self.norm = nn.LayerNorm(dim) if norm else 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,
+    ) -> tuple:
+        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
+        return self.dropout(x), weights
+
+
+class FFN(nn.Module):
+    """Feed-forward network."""
+
+    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)
+        self.activation = (
+            _get_activation_fn(activation)
+            if isinstance(activation, str)
+            else activation
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r = x
+        x = self.norm(x)
+        x = self.linear2(self.activation(self.linear1(x)))
+        return r + x
+
+
+class PixelFFN(nn.Module):
+    """Pixel-level feed-forward using CAResBlock."""
+
+    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:
+        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):
+    """Output projection FFN."""
+
+    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)
+        self.activation = (
+            _get_activation_fn(activation)
+            if isinstance(activation, str)
+            else 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
+    if activation == "gelu":
+        return F.gelu
+    raise RuntimeError("activation should be relu/gelu, not %s" % activation)

src/matanyone2/model/utils/__init__.py

@@ -0,0 +1,4 @@
+"""MatAnyone2 model utilities."""
+from . import resnet
+
+__all__ = ["resnet"]

src/matanyone2/model/utils/memory_utils.py

@@ -0,0 +1,111 @@
+"""
+Memory utilities for MatAnyone2 (similarity, softmax, affinity, readout).
+"""
+
+import math
+from typing import Optional, Tuple, Union
+
+import torch
+
+
+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:
+    """Computes similarity for memory reading."""
+    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
+
+    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
+
+    if qe is not None:
+        mk = mk.transpose(1, 2)
+        a_sq = mk.pow(2) @ qe
+        two_ab = 2 * (mk @ (qk * qe))
+        b_sq = (qe * qk.pow(2)).sum(1, keepdim=True)
+        similarity = -a_sq + two_ab - b_sq
+    else:
+        a_sq = mk.pow(2).sum(1).unsqueeze(2)
+        two_ab = 2 * (mk.transpose(1, 2) @ qk)
+        similarity = -a_sq + two_ab
+
+    if ms is not None:
+        similarity = similarity * ms / math.sqrt(CK)
+    else:
+        similarity = similarity / math.sqrt(CK)
+
+    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]]:
+    """Normalizes similarity with top-k softmax."""
+    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)
+            affinity = similarity
+        else:
+            affinity = torch.zeros_like(similarity).scatter_(1, indices, x_exp)
+    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 for get_similarity + do_softmax."""
+    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:
+    """Reads memory values using affinity."""
+    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

src/matanyone2/model/utils/parameter_groups.py

@@ -0,0 +1,66 @@
+"""
+Parameter group utilities for MatAnyone2 training.
+"""
+
+import logging
+
+log = logging.getLogger()
+
+
+def get_parameter_groups(model, stage_cfg, print_log=False):
+    """
+    Assigns different weight decays and learning rates to different parameters.
+    Returns parameter groups for 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]
+
+    memo = set()
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+        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("%s counted as a backbone parameter.", name)
+        else:
+            for e in embedding_names:
+                if name.endswith(e):
+                    embed_params.append(param)
+                    inserted = True
+                    if print_log:
+                        log.info("%s counted as an embedding parameter.", name)
+                    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

src/matanyone2/model/utils/resnet.py

@@ -0,0 +1,186 @@
+"""
+Modified ResNet structure for MatAnyone2.
+Appends extra channels to the first conv for mask input.
+"""
+
+import math
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+
+try:
+    from torch.hub import load_state_dict_from_url
+except ImportError:
+    from torch.utils.model_zoo import (
+        load_url as load_state_dict_from_url,  # type: ignore
+    )
+
+
+def load_weights_add_extra_dim(target, source_state, extra_dim=1):
+    """Loads weights, adding extra input channels if needed."""
+    new_dict = OrderedDict()
+
+    for k1, v1 in target.state_dict().items():
+        if "num_batches_tracked" not in k1:
+            if k1 in source_state:
+                tar_v = source_state[k1]
+
+                if 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": "https://download.pytorch.org/models/resnet18-5c106cde.pth",
+    "resnet50": "https://download.pytorch.org/models/resnet50-19c8e357.pth",
+}
+
+
+def conv3x3(in_planes, out_planes, stride=1, dilation=1):
+    """3x3 convolution with padding."""
+    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().__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().__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):
+    """ResNet backbone with optional extra input channels."""
+
+    def __init__(self, block, layers=(3, 4, 23, 3), extra_dim=0):
+        super().__init__()
+        self.inplanes = 64
+        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.0 / 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 _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes, dilation=dilation))
+        return nn.Sequential(*layers)
+
+
+def resnet18(pretrained=True, extra_dim=0):
+    """Build ResNet-18."""
+    model = ResNet(BasicBlock, [2, 2, 2, 2], extra_dim)
+    if pretrained:
+        state = load_state_dict_from_url(MODEL_URLS["resnet18"])
+        load_weights_add_extra_dim(model, state, extra_dim)
+    return model
+
+
+def resnet50(pretrained=True, extra_dim=0):
+    """Build ResNet-50."""
+    model = ResNet(Bottleneck, [3, 4, 6, 3], extra_dim)
+    if pretrained:
+        state = load_state_dict_from_url(MODEL_URLS["resnet50"])
+        load_weights_add_extra_dim(model, state, extra_dim)
+    return model

src/matanyone2/utils/__init__.py

@@ -0,0 +1 @@
+"""MatAnyone2 utilities."""

src/matanyone2/utils/device.py

@@ -0,0 +1,47 @@
+"""
+Device utilities for MatAnyone2 inference.
+"""
+
+import contextlib
+import functools
+
+import torch
+
+
+def get_default_device():
+    """Returns CUDA, MPS, or CPU device based on availability."""
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    elif torch.backends.mps.is_built() and torch.backends.mps.is_available():
+        return torch.device("mps")
+    else:
+        return torch.device("cpu")
+
+
+def safe_autocast_decorator(enabled=True):
+    """Decorator for safe autocast on CUDA/CPU (skips MPS)."""
+
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            device = get_default_device()
+            if device.type in ["cuda", "cpu"]:
+                with torch.amp.autocast(device_type=device.type, enabled=enabled):
+                    return func(*args, **kwargs)
+            else:
+                return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+@contextlib.contextmanager
+def safe_autocast(enabled=True):
+    """Context manager for safe autocast on CUDA/CPU (skips MPS)."""
+    device = get_default_device()
+    if device.type in ["cuda", "cpu"]:
+        with torch.amp.autocast(device_type=device.type, enabled=enabled):
+            yield
+    else:
+        yield

src/matanyone2/utils/get_default_model.py

@@ -0,0 +1,43 @@
+"""
+Loads MatAnyone2 model from checkpoint.
+"""
+
+from pathlib import Path
+
+import torch
+from omegaconf import OmegaConf, open_dict
+
+from ..model.matanyone2 import MatAnyone2
+
+
+def get_matanyone2_model(ckpt_path: str, device=None) -> MatAnyone2:
+    """Loads MatAnyone2 model from checkpoint path."""
+    config_dir = Path(__file__).resolve().parent.parent / "config"
+
+    # Load base model config and eval config
+    model_cfg = OmegaConf.load(config_dir / "model" / "base.yaml")
+    eval_cfg = OmegaConf.load(config_dir / "eval_matanyone_config.yaml")
+
+    # Clean up hydra/defaults from eval_cfg if they exist
+    if "defaults" in eval_cfg:
+        del eval_cfg["defaults"]
+    if "hydra" in eval_cfg:
+        del eval_cfg["hydra"]
+
+    # Merge them manually as hydra did
+    cfg = OmegaConf.merge(eval_cfg, {"model": model_cfg})
+    OmegaConf.resolve(cfg)
+
+    with open_dict(cfg):
+        cfg["weights"] = ckpt_path
+
+    if device is not None:
+        matanyone2 = MatAnyone2(cfg, single_object=True).to(device).eval()
+        model_weights = torch.load(cfg.weights, map_location=device)
+    else:
+        matanyone2 = MatAnyone2(cfg, single_object=True).cuda().eval()
+        model_weights = torch.load(cfg.weights)
+
+    matanyone2.load_weights(model_weights)
+
+    return matanyone2

src/matanyone2/utils/inference_utils.py

@@ -0,0 +1,69 @@
+"""
+Inference utilities for MatAnyone2 (erosion, dilation, frame loading).
+"""
+
+import os
+import random
+
+import cv2
+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):
+    """Reads frames from video file or image folder. Returns (frames, fps, length, name)."""
+    if frame_root.endswith(VIDEO_EXTENSIONS):
+        video_name = os.path.basename(frame_root)[:-4]
+        frames, _, info = torchvision.io.read_video(
+            filename=frame_root, pts_unit="sec", output_format="TCHW"
+        )
+        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]]
+            frames.append(frame)
+        fps = 24
+        frames = torch.Tensor(np.array(frames)).permute(0, 3, 1, 2).contiguous()
+
+    length = frames.shape[0]
+    return frames, fps, length, video_name
+
+
+def get_video_paths(input_root):
+    """Collects video file paths under 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):
+    """Parses comma-separated string to list of ints."""
+    return list(map(int, value.split(",")))
+
+
+def gen_dilate(alpha, min_kernel_size, max_kernel_size):
+    """Dilates alpha mask with random 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):
+    """Erodes alpha mask with random 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)

src/matanyone2/utils/tensor_utils.py

@@ -0,0 +1,71 @@
+"""
+Tensor utilities for MatAnyone2.
+"""
+
+from typing import Iterable, List
+
+import torch
+import torch.nn.functional as F
+
+from .device import safe_autocast
+
+
+def pad_divide_by(in_img: torch.Tensor, d: int) -> tuple[torch.Tensor, Iterable[int]]:
+    """Pads image so dimensions are divisible by d."""
+    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:
+    """Removes padding from image."""
+    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
+
+
+def aggregate(prob: torch.Tensor, dim: int) -> torch.Tensor:
+    """Aggregates probability tensor along dimension."""
+    with safe_autocast(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)))
+        return logits
+
+
+def cls_to_one_hot(cls_gt: torch.Tensor, num_objects: int) -> torch.Tensor:
+    """Converts class indices to one-hot encoding."""
+    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