nodes/1038lab-ComfyUI-RMBG
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add files via upload

Browse Source
This commit is contained in:
AI Lab
2025-11-24 20:59:52 -08:00
committed by GitHub
parent e34653b271
commit 0829b88776
17 changed files with 2290 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sam3/perflib/__init__.py
+8
View File
@@ -0,0 +1,8 @@
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
import os
is_enabled = False
if os.getenv("USE_PERFLIB", "1") == "1":
# print("Enabled the use of perflib.\n", end="")
is_enabled = True
sam3/perflib/associate_det_trk.py
+137
View File
@@ -0,0 +1,137 @@
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
from collections import defaultdict
import torch
import torch.nn.functional as F
from sam3.perflib.masks_ops import mask_iou
from scipy.optimize import linear_sum_assignment
def associate_det_trk(
det_masks,
track_masks,
iou_threshold=0.5,
iou_threshold_trk=0.5,
det_scores=None,
new_det_thresh=0.0,
):
"""
Optimized implementation of detection <-> track association that minimizes DtoH syncs.
Args:
det_masks: (N, H, W) tensor of predicted masks
track_masks: (M, H, W) tensor of track masks
Returns:
new_det_indices: list of indices in det_masks considered 'new'
unmatched_trk_indices: list of indices in track_masks considered 'unmatched'
"""
with torch.autograd.profiler.record_function("perflib: associate_det_trk"):
assert isinstance(det_masks, torch.Tensor), "det_masks should be a tensor"
assert isinstance(track_masks, torch.Tensor), "track_masks should be a tensor"
if det_masks.size(0) == 0 or track_masks.size(0) == 0:
return list(range(det_masks.size(0))), [], {}, {} # all detections are new
if list(det_masks.shape[-2:]) != list(track_masks.shape[-2:]):
# resize to the smaller size to save GPU memory
if torch.numel(det_masks[-2:]) < torch.numel(track_masks[-2:]):
track_masks = (
F.interpolate(
track_masks.unsqueeze(1).float(),
size=det_masks.shape[-2:],
mode="bilinear",
align_corners=False,
).squeeze(1)
> 0
)
else:
# resize detections to track size
det_masks = (
F.interpolate(
det_masks.unsqueeze(1).float(),
size=track_masks.shape[-2:],
mode="bilinear",
align_corners=False,
).squeeze(1)
> 0
)
det_masks = det_masks > 0
track_masks = track_masks > 0
iou = mask_iou(det_masks, track_masks) # (N, M)
igeit = iou >= iou_threshold
igeit_any_dim_1 = igeit.any(dim=1)
igeit_trk = iou >= iou_threshold_trk
iou_list = iou.cpu().numpy().tolist()
igeit_list = igeit.cpu().numpy().tolist()
igeit_any_dim_1_list = igeit_any_dim_1.cpu().numpy().tolist()
igeit_trk_list = igeit_trk.cpu().numpy().tolist()
det_scores_list = (
det_scores
if det_scores is None
else det_scores.cpu().float().numpy().tolist()
)
# Hungarian matching for tracks (one-to-one: each track matches at most one detection)
# For detections: allow many tracks to match to the same detection (many-to-one)
# If either is empty, return all detections as new
if det_masks.size(0) == 0 or track_masks.size(0) == 0:
return list(range(det_masks.size(0))), [], {}
# Hungarian matching: maximize IoU for tracks
cost_matrix = 1 - iou.cpu().numpy() # Hungarian solves for minimum cost
row_ind, col_ind = linear_sum_assignment(cost_matrix)
def branchy_hungarian_better_uses_the_cpu(
cost_matrix, row_ind, col_ind, iou_list, det_masks, track_masks
):
matched_trk = set()
matched_det = set()
matched_det_scores = {} # track index -> [det_score, det_score * iou] det score of matched detection mask
for d, t in zip(row_ind, col_ind):
matched_det_scores[t] = [
det_scores_list[d],
det_scores_list[d] * iou_list[d][t],
]
if igeit_trk_list[d][t]:
matched_trk.add(t)
matched_det.add(d)
# Tracks not matched by Hungarian assignment above threshold are unmatched
unmatched_trk_indices = [
t for t in range(track_masks.size(0)) if t not in matched_trk
]
# For detections: allow many tracks to match to the same detection (many-to-one)
# So, a detection is 'new' if it does not match any track above threshold
assert track_masks.size(0) == igeit.size(
1
) # Needed for loop optimizaiton below
new_det_indices = []
for d in range(det_masks.size(0)):
if not igeit_any_dim_1_list[d]:
if det_scores is not None and det_scores[d] >= new_det_thresh:
new_det_indices.append(d)
# for each detection, which tracks it matched to (above threshold)
det_to_matched_trk = defaultdict(list)
for d in range(det_masks.size(0)):
for t in range(track_masks.size(0)):
if igeit_list[d][t]:
det_to_matched_trk[d].append(t)
return (
new_det_indices,
unmatched_trk_indices,
det_to_matched_trk,
matched_det_scores,
)
return (branchy_hungarian_better_uses_the_cpu)(
cost_matrix, row_ind, col_ind, iou_list, det_masks, track_masks
)
sam3/perflib/compile.py
+99
View File
@@ -0,0 +1,99 @@
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
import torch
def recursive_fn_factory(fn):
def recursive_fn(b):
if isinstance(b, dict):
return {k: recursive_fn(b[k]) for k in b}
if isinstance(b, list):
return [recursive_fn(t) for t in b]
if isinstance(b, tuple):
return tuple(recursive_fn(t) for t in b)
if isinstance(b, torch.Tensor):
return fn(b)
# Yes, writing out an explicit white list of
# trivial types is tedious, but so are bugs that
# come from not applying fn, when expected to have
# applied it.
if b is None:
return b
trivial_types = [bool, int]
for t in trivial_types:
if isinstance(b, t):
return b
raise TypeError(f"Unexpected type {type(b)}")
return recursive_fn
recursive_contiguous = recursive_fn_factory(lambda x: x.contiguous())
recursive_clone = recursive_fn_factory(torch.clone)
def compile_wrapper(
fn, *, mode="max-autotune", fullgraph=True, dynamic=False, name=None
):
compiled_fn = torch.compile(fn, mode=mode, fullgraph=fullgraph, dynamic=dynamic)
def compiled_fn_wrapper(*args, **kwargs):
with torch.autograd.profiler.record_function(
f"compiled {fn}" if name is None else name
):
cont_args = recursive_contiguous(args)
cont_kwargs = recursive_contiguous(kwargs)
result = compiled_fn(*cont_args, **cont_kwargs)
cloned_result = recursive_clone(result)
return cloned_result
return compiled_fn_wrapper
def shape_logging_wrapper(fn, keep_kwargs, enable_logging=False):
"""
Wraps a function and prints the shapes of all tensor inputs.
Only prints when a new combination of shapes is seen.
Thread-safe.
Args:
fn: Function to wrap
enable_logging: Boolean flag to enable/disable logging
"""
seen_shapes = set()
def get_shape(obj):
if isinstance(obj, torch.Tensor):
return obj.shape
elif isinstance(obj, (list, tuple)):
if len(obj) > 1:
return tuple(get_shape(x) for x in obj)
return get_shape(obj[0])
elif isinstance(obj, dict):
return tuple(sorted((k, get_shape(v)) for k, v in obj.items()))
else:
return type(obj).__name__
def wrapper(*args, **kwargs):
shapes = tuple(get_shape(arg) for arg in args) + tuple(
(k, get_shape(v))
for k, v in kwargs.items()
if isinstance(v, (torch.Tensor, list))
and (len(keep_kwargs) > 0 and k in keep_kwargs)
)
if shapes not in seen_shapes:
seen_shapes.add(shapes)
if enable_logging:
print(f"[ShapeLogger] New input shapes for {fn.__qualname__}: {shapes}")
return fn(*args, **kwargs)
# Allow toggling the flag at runtime
wrapper.enable_logging = enable_logging
def set_logging(enabled=False):
nonlocal enable_logging
enable_logging = enabled
wrapper.enable_logging = enable_logging
wrapper.set_logging = set_logging
return wrapper
sam3/perflib/connected_components.py
+84
View File
@@ -0,0 +1,84 @@
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
import logging
import torch
try:
from cc_torch import get_connected_components
HAS_CC_TORCH = True
except ImportError:
logging.debug(
"cc_torch not found. Consider installing for better performance. Command line:"
" pip install git+https://github.com/ronghanghu/cc_torch.git"
)
HAS_CC_TORCH = False
def connected_components_cpu_single(values: torch.Tensor):
assert values.dim() == 2
from skimage.measure import label
labels, num = label(values.cpu().numpy(), return_num=True)
labels = torch.from_numpy(labels)
counts = torch.zeros_like(labels)
for i in range(1, num + 1):
cur_mask = labels == i
cur_count = cur_mask.sum()
counts[cur_mask] = cur_count
return labels, counts
def connected_components_cpu(input_tensor: torch.Tensor):
out_shape = input_tensor.shape
if input_tensor.dim() == 4 and input_tensor.shape[1] == 1:
input_tensor = input_tensor.squeeze(1)
else:
assert (
input_tensor.dim() == 3
), "Input tensor must be (B, H, W) or (B, 1, H, W)."
batch_size = input_tensor.shape[0]
labels_list = []
counts_list = []
for b in range(batch_size):
labels, counts = connected_components_cpu_single(input_tensor[b])
labels_list.append(labels)
counts_list.append(counts)
labels_tensor = torch.stack(labels_list, dim=0).to(input_tensor.device)
counts_tensor = torch.stack(counts_list, dim=0).to(input_tensor.device)
return labels_tensor.view(out_shape), counts_tensor.view(out_shape)
def connected_components(input_tensor: torch.Tensor):
"""
Computes connected components labeling on a batch of 2D tensors, using the best available backend.
Args:
input_tensor (torch.Tensor): A BxHxW integer tensor or Bx1xHxW. Non-zero values are considered foreground. Bool tensor also accepted
Returns:
Tuple[torch.Tensor, torch.Tensor]: Both tensors have the same shape as input_tensor.
- A tensor with dense labels. Background is 0.
- A tensor with the size of the connected component for each pixel.
"""
if input_tensor.dim() == 3:
input_tensor = input_tensor.unsqueeze(1)
assert (
input_tensor.dim() == 4 and input_tensor.shape[1] == 1
), "Input tensor must be (B, H, W) or (B, 1, H, W)."
if input_tensor.is_cuda:
if HAS_CC_TORCH:
return get_connected_components(input_tensor.to(torch.uint8))
else:
# triton fallback
from sam3.perflib.triton.connected_components import (
connected_components_triton,
)
return connected_components_triton(input_tensor)
# CPU fallback
return connected_components_cpu(input_tensor)
sam3/perflib/fa3.py
+27
View File
@@ -0,0 +1,27 @@
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
import torch
@torch.library.custom_op("flash::flash_attn_func", mutates_args=())
def flash_attn_func_op(
q: torch.Tensor, k: torch.Tensor, v: torch.Tensor
) -> torch.Tensor:
from flash_attn_interface import flash_attn_func as fa3
return fa3(q, k, v)
def flash_attn_func(q, k, v):
dtype = torch.float8_e4m3fn
return flash_attn_func_op(q.to(dtype), k.to(dtype), v.to(dtype)).to(q.dtype)
@flash_attn_func_op.register_fake
def _(q, k, v, **kwargs):
# two outputs:
# 1. output: (batch, seq_len, num_heads, head_dim)
# 2. softmax_lse: (batch, num_heads, seq_len) with dtype=torch.float32
# output needs to be bfloat16, not float8!
meta_q = torch.empty_like(q, dtype=torch.bfloat16).contiguous()
return meta_q
sam3/perflib/masks_ops.py
+69
View File
@@ -0,0 +1,69 @@
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
import torch
def masks_to_boxes(masks: torch.Tensor, obj_ids: list[int]):
with torch.autograd.profiler.record_function("perflib: masks_to_boxes"):
# Sanity check based on callsite for replacement
assert masks.shape[0] == len(obj_ids)
assert masks.dim() == 3
# Based on torchvision masks_to_boxes
if masks.numel() == 0:
return torch.zeros((0, 4), device=masks.device, dtype=torch.float)
N, H, W = masks.shape
device = masks.device
y = torch.arange(H, device=device).view(1, H)
x = torch.arange(W, device=device).view(1, W)
masks_with_obj = masks != 0 # N, H, W
masks_with_obj_x = masks_with_obj.amax(
dim=1
) # N, H (which columns have objects)
masks_with_obj_y = masks_with_obj.amax(dim=2) # N, W (which rows have objects)
masks_without_obj_x = ~masks_with_obj_x
masks_without_obj_y = ~masks_with_obj_y
bounding_boxes_0 = torch.amin(
(masks_without_obj_x * W) + (masks_with_obj_x * x), dim=1
)
bounding_boxes_1 = torch.amin(
(masks_without_obj_y * H) + (masks_with_obj_y * y), dim=1
)
bounding_boxes_2 = torch.amax(masks_with_obj_x * x, dim=1)
bounding_boxes_3 = torch.amax(masks_with_obj_y * y, dim=1)
bounding_boxes = torch.stack(
[bounding_boxes_0, bounding_boxes_1, bounding_boxes_2, bounding_boxes_3],
dim=1,
).to(dtype=torch.float)
assert bounding_boxes.shape == (N, 4)
assert bounding_boxes.device == masks.device
assert bounding_boxes.dtype == torch.float
return bounding_boxes
def mask_iou(pred_masks: torch.Tensor, gt_masks: torch.Tensor) -> torch.Tensor:
"""
Compute the IoU (Intersection over Union) between predicted masks and ground truth masks.
Args:
- pred_masks: (N, H, W) bool Tensor, containing binary predicted segmentation masks
- gt_masks: (M, H, W) bool Tensor, containing binary ground truth segmentation masks
Returns:
- ious: (N, M) float Tensor, containing IoUs for each pair of predicted and ground truth masks
"""
assert pred_masks.dtype == gt_masks.dtype == torch.bool
N, H, W = pred_masks.shape
M, _, _ = gt_masks.shape
# Flatten masks: (N, 1, H*W) and (1, M, H*W)
pred_flat = pred_masks.view(N, 1, H * W)
gt_flat = gt_masks.view(1, M, H * W)
# Compute intersection and union: (N, M)
intersection = (pred_flat & gt_flat).sum(dim=2).float()
union = (pred_flat | gt_flat).sum(dim=2).float()
ious = intersection / union.clamp(min=1)
return ious # shape: (N, M)
sam3/perflib/nms.py
+91
View File
@@ -0,0 +1,91 @@
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
import logging
import numpy as np
import torch
from sam3.perflib.masks_ops import mask_iou
try:
from torch_generic_nms import generic_nms as generic_nms_cuda
GENERIC_NMS_AVAILABLE = True
except ImportError:
logging.debug(
"Falling back to triton or CPU mask NMS implementation -- please install `torch_generic_nms` via\n\t"
'pip uninstall -y torch_generic_nms; TORCH_CUDA_ARCH_LIST="8.0 9.0" pip install git+https://github.com/ronghanghu/torch_generic_nms'
)
GENERIC_NMS_AVAILABLE = False
def nms_masks(
pred_probs: torch.Tensor,
pred_masks: torch.Tensor,
prob_threshold: float,
iou_threshold: float,
) -> torch.Tensor:
"""
Args:
- pred_probs: (num_det,) float Tensor, containing the score (probability) of each detection
- pred_masks: (num_det, H_mask, W_mask) float Tensor, containing the binary segmentation mask of each detection
- prob_threshold: float, score threshold to prefilter detections (NMS is performed on detections above threshold)
- iou_threshold: float, mask IoU threshold for NMS
Returns:
- keep: (num_det,) bool Tensor, indicating whether each detection is kept after score thresholding + NMS
"""
# prefilter the detections with prob_threshold ("valid" are those above prob_threshold)
is_valid = pred_probs > prob_threshold # (num_det,)
probs = pred_probs[is_valid] # (num_valid,)
masks_binary = pred_masks[is_valid] > 0 # (num_valid, H_mask, W_mask)
if probs.numel() == 0:
return is_valid # no valid detection, return empty keep mask
ious = mask_iou(masks_binary, masks_binary) # (num_valid, num_valid)
kept_inds = generic_nms(ious, probs, iou_threshold)
# valid_inds are the indices among `probs` of valid detections before NMS (or -1 for invalid)
valid_inds = torch.where(is_valid, is_valid.cumsum(dim=0) - 1, -1) # (num_det,)
keep = torch.isin(valid_inds, kept_inds) # (num_det,)
return keep
def generic_nms(
ious: torch.Tensor, scores: torch.Tensor, iou_threshold=0.5
) -> torch.Tensor:
"""A generic version of `torchvision.ops.nms` that takes a pairwise IoU matrix."""
assert ious.dim() == 2 and ious.size(0) == ious.size(1)
assert scores.dim() == 1 and scores.size(0) == ious.size(0)
if ious.is_cuda:
if GENERIC_NMS_AVAILABLE:
return generic_nms_cuda(ious, scores, iou_threshold, use_iou_matrix=True)
else:
from sam3.perflib.triton.nms import nms_triton
return nms_triton(ious, scores, iou_threshold)
return generic_nms_cpu(ious, scores, iou_threshold)
def generic_nms_cpu(
ious: torch.Tensor, scores: torch.Tensor, iou_threshold=0.5
) -> torch.Tensor:
"""
A generic version of `torchvision.ops.nms` that takes a pairwise IoU matrix. (CPU implementation
based on https://github.com/jwyang/faster-rcnn.pytorch/blob/master/lib/model/nms/nms_cpu.py)
"""
ious_np = ious.float().detach().cpu().numpy()
scores_np = scores.float().detach().cpu().numpy()
order = scores_np.argsort()[::-1]
kept_inds = []
while order.size > 0:
i = order.item(0)
kept_inds.append(i)
inds = np.where(ious_np[i, order[1:]] <= iou_threshold)[0]
order = order[inds + 1]
return torch.tensor(kept_inds, dtype=torch.int64, device=scores.device)
sam3/perflib/tests/assets/masks.tiff
BIN
View File
Binary file not shown.
sam3/perflib/tests/tests.py
+59
View File
@@ -0,0 +1,59 @@
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
import os
import numpy as np
import pytest
import torch
from PIL import Image
from sam3.perflib.masks_ops import masks_to_boxes
class TestMasksToBoxes:
def test_masks_box(self):
def masks_box_check(masks, expected, atol=1e-4):
out = masks_to_boxes(masks, [1 for _ in range(masks.shape[0])])
assert out.dtype == torch.float
print("out: ", out)
print("expected: ", expected)
torch.testing.assert_close(
out, expected, rtol=0.0, check_dtype=True, atol=atol
)
# Check for int type boxes.
def _get_image():
assets_directory = os.path.join(
os.path.dirname(os.path.abspath(__file__)), "assets"
)
mask_path = os.path.join(assets_directory, "masks.tiff")
image = Image.open(mask_path)
return image
def _create_masks(image, masks):
for index in range(image.n_frames):
image.seek(index)
frame = np.array(image)
masks[index] = torch.tensor(frame)
return masks
expected = torch.tensor(
[
[127, 2, 165, 40],
[2, 50, 44, 92],
[56, 63, 98, 100],
[139, 68, 175, 104],
[160, 112, 198, 145],
[49, 138, 99, 182],
[108, 148, 152, 213],
],
dtype=torch.float,
)
image = _get_image()
for dtype in [torch.float16, torch.float32, torch.float64]:
masks = torch.zeros(
(image.n_frames, image.height, image.width), dtype=dtype
)
masks = _create_masks(image, masks)
masks_box_check(masks, expected)
sam3/perflib/triton/connected_components.py
+468
View File
@@ -0,0 +1,468 @@
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
import math
import torch
import triton
import triton.language as tl
@triton.jit
def _any_combine(a, b):
return a | b
@triton.jit
def tl_any(a, dim=0):
return tl.reduce(a, dim, _any_combine)
# ==============================================================================
# ## Phase 1: Initialization Kernel
# ==============================================================================
# Each foreground pixel (value > 0) gets a unique label equal to its
# linear index. Background pixels (value == 0) get a sentinel label of -1.
# Note that the indexing is done across batch boundaries for simplicity
# (i.e., the first pixel of image 1 gets label H*W, etc.)
@triton.jit
def _init_labels_kernel(
input_ptr, labels_ptr, numel: tl.constexpr, BLOCK_SIZE: tl.constexpr
):
pid = tl.program_id(0)
offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = offsets < numel
input_values = tl.load(input_ptr + offsets, mask=mask, other=0)
indices = tl.where((input_values != 0), offsets, -1)
tl.store(labels_ptr + offsets, indices, mask=mask)
# ==============================================================================
# ## Phase 2: Local merging
# ==============================================================================
# Each pixel tries to merge with its 8-connected neighbors (up, down, left, right)
# if they have the same value. This is done using a disjoint-set union operation.
@triton.jit
def find(labels_ptr, indices, mask):
current_pids = indices
# 'is_done' tracks lanes that have finished their work.
# A lane is initially "done" if it's not active (mask is False).
is_done = ~mask
# Loop as long as there is at least one lane that is NOT done.
while tl_any(~is_done):
# The work_mask is for lanes that are still active and seeking their root.
work_mask = ~is_done
parents = tl.load(labels_ptr + current_pids, mask=work_mask, other=-1)
# A lane is now done if its parent is itself (it's a root)
# or if it hits a -1 sentinel (a safe exit condition).
is_root = parents == current_pids
is_sentinel = parents == -1
is_done |= is_root | is_sentinel
# For lanes that are not yet done, update their pid to their parent to continue traversal.
current_pids = tl.where(is_done, current_pids, parents)
# We could add the following line to do path compression, but experimentally it's slower
# tl.atomic_min(labels_ptr + indices, current_pids, mask=mask)
return current_pids
@triton.jit
def union(labels_ptr, a, b, process_mask):
# This function implements a disjoint-set union
# As an invariant, we use the fact that the roots have the lower id. That helps parallelization
# However, that is not sufficient by itself. Suppose two threads want to do union(0,2) and union(1,2) at the same time
# Then if we do a naive atomic_min, 0 and 1 will compete to be the new parent of 2 and min(0, 1) will win.
# However, 1 still needs to be merged with the new {0, 2} component.
# To ensure that merge is also done, we need to detect whether the merge was successful, and if not retry until it is
current_a = a
current_b = b
final_root = a
# A mask to track which lanes have successfully completed their union.
done_mask = ~process_mask # tl.zeros_like(a) == 1 # Init with all False
while tl_any(~done_mask):
# Define the mask for lanes that still need work in this iteration
work_mask = process_mask & ~done_mask
# Find the roots for the current a and b values in the active lanes
root_a = find(labels_ptr, current_a, work_mask)
tl.debug_barrier()
root_b = find(labels_ptr, current_b, work_mask)
# 7. Merge logic
# If roots are already the same, the sets are already merged. Mark as done.
are_equal = root_a == root_b
final_root = tl.where(are_equal & work_mask & ~done_mask, root_a, final_root)
done_mask |= are_equal & work_mask
# Define masks for the two merge cases (a < b or b < a)
a_is_smaller = root_a < root_b
# Case 1: root_a < root_b. Attempt to set parent[root_b] = root_a
merge_mask_a_smaller = work_mask & a_is_smaller & ~are_equal
ptr_b = labels_ptr + root_b
old_val_b = tl.atomic_min(ptr_b, root_a, mask=merge_mask_a_smaller)
# A lane is done if its atomic op was successful (old value was what we expected)
success_b = old_val_b == root_b
final_root = tl.where(success_b & work_mask & ~done_mask, root_a, final_root)
done_mask |= success_b & merge_mask_a_smaller
# *** Crucial Retry Logic ***
# If the update failed (old_val_b != root_b), another thread interfered.
# We update `current_b` to this new root (`old_val_b`) and will retry in the next loop iteration.
current_b = tl.where(success_b | ~merge_mask_a_smaller, current_b, old_val_b)
# Case 2: root_b < root_a. Attempt to set parent[root_a] = root_b
merge_mask_b_smaller = work_mask & ~a_is_smaller & ~are_equal
ptr_a = labels_ptr + root_a
old_val_a = tl.atomic_min(ptr_a, root_b, mask=merge_mask_b_smaller)
success_a = old_val_a == root_a
final_root = tl.where(success_a & work_mask & ~done_mask, root_b, final_root)
done_mask |= success_a & merge_mask_b_smaller
# *** Crucial Retry Logic ***
# Similarly, update `current_a` if the atomic operation failed.
current_a = tl.where(success_a | ~merge_mask_b_smaller, current_a, old_val_a)
return final_root
@triton.jit
def _merge_helper(
input_ptr,
labels_ptr,
base_offset,
offsets_h,
offsets_w,
mask_2d,
valid_current,
current_values,
current_labels,
H,
W,
dx: tl.constexpr,
dy: tl.constexpr,
):
# Helper functions to compute merge with a specific neighbor offset (dx, dy)
neighbor_h = offsets_h + dy
neighbor_w = offsets_w + dx
# Proper bounds checking: all four bounds must be satisfied
mask_n = (
mask_2d
& (neighbor_h[:, None] >= 0)
& (neighbor_h[:, None] < H)
& (neighbor_w[None, :] >= 0)
& (neighbor_w[None, :] < W)
)
offsets_neighbor = neighbor_h[:, None] * W + neighbor_w[None, :]
neighbor_values = tl.load(
input_ptr + base_offset + offsets_neighbor, mask=mask_n, other=-1
)
mask_n = tl.ravel(mask_n)
neighbor_labels = tl.load(
labels_ptr + tl.ravel(base_offset + offsets_neighbor), mask=mask_n, other=-1
)
to_merge = (
mask_n & (neighbor_labels != -1) & tl.ravel(current_values == neighbor_values)
)
valid_write = valid_current & to_merge
# returns new parents for the pixels that were merged (otherwise keeps current labels)
return tl.where(
valid_write,
union(labels_ptr, current_labels, neighbor_labels, valid_write),
current_labels,
)
@triton.autotune(
configs=[
triton.Config(
{"BLOCK_SIZE_H": 4, "BLOCK_SIZE_W": 16}, num_stages=1, num_warps=2
),
triton.Config(
{"BLOCK_SIZE_H": 4, "BLOCK_SIZE_W": 32}, num_stages=2, num_warps=4
),
],
key=["H", "W"],
restore_value=["labels_ptr"],
)
@triton.jit
def _local_prop_kernel(
labels_ptr,
input_ptr,
H: tl.constexpr,
W: tl.constexpr,
BLOCK_SIZE_H: tl.constexpr,
BLOCK_SIZE_W: tl.constexpr,
):
# This is the meat of the Phase 2 to do local merging
# It will be launched with a 2D grid:
# - dim 0: batch index
# - dim 1: block index over HxW image (2D tiling)
pid_b = tl.program_id(0)
pid_hw = tl.program_id(1)
# Calculate offsets for the core block
offsets_h = (pid_hw // tl.cdiv(W, BLOCK_SIZE_W)) * BLOCK_SIZE_H + tl.arange(
0, BLOCK_SIZE_H
)
offsets_w = (pid_hw % tl.cdiv(W, BLOCK_SIZE_W)) * BLOCK_SIZE_W + tl.arange(
0, BLOCK_SIZE_W
)
base_offset = pid_b * H * W
offsets_2d = offsets_h[:, None] * W + offsets_w[None, :]
mask_2d = (offsets_h[:, None] < H) & (offsets_w[None, :] < W)
mask_1d = tl.ravel(mask_2d)
# Load the current labels for the block - these are parent pointers
current_labels = tl.load(
labels_ptr + tl.ravel(base_offset + offsets_2d), mask=mask_1d, other=-1
)
current_values = tl.load(
input_ptr + base_offset + offsets_2d, mask=mask_2d, other=-1
)
valid_current = mask_1d & (current_labels != -1)
# Horizontal merge
current_labels = _merge_helper(
input_ptr,
labels_ptr,
base_offset,
offsets_h,
offsets_w,
mask_2d,
valid_current,
current_values,
current_labels,
H,
W,
-1,
0,
)
# Vertical merge
current_labels = _merge_helper(
input_ptr,
labels_ptr,
base_offset,
offsets_h,
offsets_w,
mask_2d,
valid_current,
current_values,
current_labels,
H,
W,
0,
-1,
)
# Diagonal merges
current_labels = _merge_helper(
input_ptr,
labels_ptr,
base_offset,
offsets_h,
offsets_w,
mask_2d,
valid_current,
current_values,
current_labels,
H,
W,
-1,
-1,
)
current_labels = _merge_helper(
input_ptr,
labels_ptr,
base_offset,
offsets_h,
offsets_w,
mask_2d,
valid_current,
current_values,
current_labels,
H,
W,
-1,
1,
)
# This actually does some path compression, in a lightweight but beneficial way
tl.atomic_min(
labels_ptr + tl.ravel(base_offset + offsets_2d), current_labels, mask=mask_1d
)
# ==============================================================================
# ## Phase 3: Pointer Jumping Kernel
# ==============================================================================
# This kernel performs pointer jumping to ensure that all pixels point directly to their root labels.
# This is done in a loop until convergence.
@triton.jit
def _pointer_jump_kernel(
labels_in_ptr, labels_out_ptr, numel: tl.constexpr, BLOCK_SIZE: tl.constexpr
):
"""
Pointer jumping kernel with double buffering to avoid race conditions.
Reads from labels_in_ptr and writes to labels_out_ptr.
"""
# This kernel is launched with a 1D grid, and does not care about batching explicitly.
# By construction, the labels are global indices across the batch, and we never perform
# cross-batch merges, so this is safe.
pid = tl.program_id(0)
offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = offsets < numel
# Load current labels from input buffer
current_labels = tl.load(labels_in_ptr + offsets, mask=mask, other=-1)
valid_mask = mask & (current_labels != -1)
# A mask to track which lanes have successfully completed their union.
done_mask = ~valid_mask
while tl_any(~(done_mask | ~valid_mask)):
parent_labels = tl.load(
labels_in_ptr + current_labels, mask=valid_mask, other=-1
)
are_equal = current_labels == parent_labels
done_mask |= are_equal & valid_mask
current_labels = tl.where(
~done_mask, tl.minimum(current_labels, parent_labels), current_labels
)
# Write to output buffer (safe because we're not reading from it)
tl.store(labels_out_ptr + offsets, current_labels, mask=mask)
# ==============================================================================
# ## Phase 4: Kernels for Computing Component Sizes
# ==============================================================================
# Step 4.1: Count occurrences of each root label using atomic adds.
@triton.jit
def _count_labels_kernel(labels_ptr, sizes_ptr, numel, BLOCK_SIZE: tl.constexpr):
pid = tl.program_id(0)
offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = offsets < numel
# Load the final, converged labels
labels = tl.load(labels_ptr + offsets, mask=mask, other=-1)
valid_mask = mask & (labels != -1)
# Atomically increment the counter for each label. This builds a histogram.
tl.atomic_add(sizes_ptr + labels, 1, mask=valid_mask)
# Step 4.2: Broadcast the computed sizes back to the output tensor.
@triton.jit
def _broadcast_sizes_kernel(
labels_ptr, sizes_ptr, out_ptr, numel, BLOCK_SIZE: tl.constexpr
):
pid = tl.program_id(0)
offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = offsets < numel
# Load the final labels
labels = tl.load(labels_ptr + offsets, mask=mask, other=-1)
valid_mask = mask & (labels != -1)
# Look up the size for each label from the histogram
component_sizes = tl.load(sizes_ptr + labels, mask=valid_mask, other=0)
# Write the size to the final output tensor. Background pixels get size 0.
tl.store(out_ptr + offsets, component_sizes, mask=mask)
def connected_components_triton(input_tensor: torch.Tensor):
"""
Computes connected components labeling on a batch of 2D integer tensors using Triton.
Args:
input_tensor (torch.Tensor): A BxHxW integer tensor or Bx1xHxW. Non-zero values are considered foreground. Bool tensor also accepted
Returns:
Tuple[torch.Tensor, int]: A tuple containing:
- A BxHxW output tensor with dense labels. Background is 0.
- A BxHxW tensor with the size of the connected component for each pixel.
"""
assert (
input_tensor.is_cuda and input_tensor.is_contiguous()
), "Input tensor must be a contiguous CUDA tensor."
out_shape = input_tensor.shape
if input_tensor.dim() == 4 and input_tensor.shape[1] == 1:
input_tensor = input_tensor.squeeze(1)
else:
assert (
input_tensor.dim() == 3
), "Input tensor must be (B, H, W) or (B, 1, H, W)."
B, H, W = input_tensor.shape
numel = B * H * W
device = input_tensor.device
# --- Allocate Tensors ---
labels = torch.empty_like(input_tensor, dtype=torch.int32)
output = torch.empty_like(input_tensor, dtype=torch.int32)
# --- Phase 1 ---
BLOCK_SIZE = 256
grid_init = (triton.cdiv(numel, BLOCK_SIZE),)
_init_labels_kernel[grid_init](
input_tensor,
labels,
numel,
BLOCK_SIZE=BLOCK_SIZE,
)
# --- Phase 2 ---
grid_local_prop = lambda meta: (
B,
triton.cdiv(H, meta["BLOCK_SIZE_H"]) * triton.cdiv(W, meta["BLOCK_SIZE_W"]),
)
_local_prop_kernel[grid_local_prop](labels, input_tensor, H, W)
# --- Phase 3 ---
BLOCK_SIZE = 256
grid_jump = lambda meta: (triton.cdiv(numel, meta["BLOCK_SIZE"]),)
_pointer_jump_kernel[grid_jump](labels, output, numel, BLOCK_SIZE=BLOCK_SIZE)
# --- Phase 4 ---
# Allocate tensor to store the final output sizes
component_sizes_out = torch.empty_like(input_tensor, dtype=torch.int32)
# Allocate a temporary 1D tensor to act as the histogram
# Size is numel because labels can be up to numel-1
sizes_histogram = torch.zeros(numel, dtype=torch.int32, device=device)
# 4.1: Count the occurrences of each label
grid_count = (triton.cdiv(numel, BLOCK_SIZE),)
_count_labels_kernel[grid_count](
output, sizes_histogram, numel, BLOCK_SIZE=BLOCK_SIZE
)
# 2.2: Broadcast the counts to the final output tensor
grid_broadcast = (triton.cdiv(numel, BLOCK_SIZE),)
_broadcast_sizes_kernel[grid_broadcast](
output, sizes_histogram, component_sizes_out, numel, BLOCK_SIZE=BLOCK_SIZE
)
return output.view(out_shape) + 1, component_sizes_out.view(out_shape)
sam3/perflib/triton/nms.py
+124
View File
@@ -0,0 +1,124 @@
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
# Adapted from https://github.com/stackav-oss/conch/blob/main/conch/kernels/vision/nms.py
import torch
import triton
import triton.language as tl
@triton.autotune(
configs=[
triton.Config({"cxpr_block_size": 128}),
triton.Config({"cxpr_block_size": 256}),
triton.Config({"cxpr_block_size": 512}),
triton.Config({"cxpr_block_size": 1024}),
triton.Config({"cxpr_block_size": 2048}),
triton.Config({"cxpr_block_size": 4096}),
triton.Config({"cxpr_block_size": 8192}),
],
key=["num_boxes"],
)
@triton.jit
def _nms_suppression_kernel(
# Tensors
iou_mask_ptr: tl.tensor, # [N, N]
keep_mask_ptr: tl.tensor, # [N]
# Scalars
num_boxes: tl.int32,
# Strides
iou_mask_stride: tl.int32,
# Constexprs
cxpr_block_size: tl.constexpr,
) -> None:
"""NMS suppression kernel.
Args:
iou_mask_ptr: Pointer to precomputed IoU mask, shape: (N, N).
keep_mask_ptr: Pointer to keep mask tensor, shape: (N,).
num_boxes: Number of boxes.
iou_mask_stride: Stride for IoU mask tensor.
cxpr_block_size: Block size for processing.
"""
# Sequential NMS: for each box in sorted order, suppress later boxes
for current_box_idx in range(num_boxes - 1):
# Check if current box is still kept
is_kept = tl.load(keep_mask_ptr + current_box_idx)
if is_kept:
# IoU mask row offset for the current box
# Because the IoU mask is sorted by score, we will only consider boxes that come after the current box.
# This means we only need to read the upper triangular part of the IoU mask.
iou_row_offset = current_box_idx * iou_mask_stride
# Only process boxes that come after the current box
next_box_idx = current_box_idx + 1
remaining_boxes = num_boxes - next_box_idx
# Iterate blockwise through the columns
for block_idx in range(tl.cdiv(remaining_boxes, cxpr_block_size)):
# Masked load of indices for the target boxes in the current block
block_start = next_box_idx + block_idx * cxpr_block_size
target_box_offsets = block_start + tl.arange(0, cxpr_block_size)
target_box_mask = target_box_offsets < num_boxes
# Suppress boxes with lower scores that have high IoU
suppression_mask = tl.load(
iou_mask_ptr + iou_row_offset + target_box_offsets,
mask=target_box_mask,
other=False,
)
suppression_mask = tl.cast(suppression_mask, tl.int1)
# Conditionally store suppression result for high-IoU boxes
tl.store(
keep_mask_ptr + target_box_offsets, False, mask=suppression_mask
)
# Potential race condition: we need to ensure all threads complete the store before the next
# iteration otherwise we may load stale data for whether or not a box has been suppressed.
tl.debug_barrier()
def nms_triton(
ious: torch.Tensor,
scores: torch.Tensor,
iou_threshold: float,
) -> torch.Tensor:
"""Perform NMS given the iou matrix, the scores and the iou threshold
Args:
ious: Pairwise IoU tensor of shape (N, N).
scores: Scores tensor of shape (N,).
iou_threshold: IoU threshold for suppression.
Returns:
Tensor: Indices of kept boxes, sorted by decreasing score.
"""
assert scores.dim() == 1, "Scores must be 1D"
iou_mask = ious > iou_threshold
assert iou_mask.dim() == 2
assert iou_mask.shape[0] == iou_mask.shape[1] == scores.shape[0]
assert iou_mask.device == scores.device
assert iou_mask.dtype == torch.bool
num_boxes = scores.size(0)
keep_mask = torch.ones(len(scores), device=scores.device, dtype=torch.bool)
# Sort boxes by scores in descending order
_, sorted_indices = torch.sort(scores, dim=0, stable=True, descending=True)
iou_mask = iou_mask[sorted_indices][:, sorted_indices].contiguous()
# For the suppression stage, we need to process sequentially, but we'll still take
# advantage of parallelism by processing in blocks in one program.
stage2_grid = (1,)
_nms_suppression_kernel[stage2_grid](
# Tensors
iou_mask_ptr=iou_mask,
keep_mask_ptr=keep_mask,
# Scalars
num_boxes=num_boxes,
# Strides
iou_mask_stride=iou_mask.stride(0),
)
# Extract indices of kept boxes
return sorted_indices[keep_mask]
sam3/sam/__init__.py
+4
View File
@@ -0,0 +1,4 @@
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
from .mask_decoder import MaskDecoder
from .prompt_encoder import PromptEncoder
from .transformer import TwoWayTransformer
sam3/sam/common.py
+39
View File
@@ -0,0 +1,39 @@
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
from typing import Type
import torch
import torch.nn as nn
class MLPBlock(nn.Module):
def __init__(
self,
embedding_dim: int,
mlp_dim: int,
act: Type[nn.Module] = nn.GELU,
) -> None:
super().__init__()
self.lin1 = nn.Linear(embedding_dim, mlp_dim)
self.lin2 = nn.Linear(mlp_dim, embedding_dim)
self.act = act()
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.lin2(self.act(self.lin1(x)))
# From https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/batch_norm.py # noqa
# Itself from https://github.com/facebookresearch/ConvNeXt/blob/d1fa8f6fef0a165b27399986cc2bdacc92777e40/models/convnext.py#L119 # noqa
class LayerNorm2d(nn.Module):
def __init__(self, num_channels: int, eps: float = 1e-6) -> None:
super().__init__()
self.weight = nn.Parameter(torch.ones(num_channels))
self.bias = nn.Parameter(torch.zeros(num_channels))
self.eps = eps
def forward(self, x: torch.Tensor) -> torch.Tensor:
u = x.mean(1, keepdim=True)
s = (x - u).pow(2).mean(1, keepdim=True)
x = (x - u) / torch.sqrt(s + self.eps)
x = self.weight[:, None, None] * x + self.bias[:, None, None]
return x
sam3/sam/mask_decoder.py
+319
View File
@@ -0,0 +1,319 @@
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
from typing import List, Optional, Tuple, Type
import torch
from torch import nn
from torch.nn import functional as F
from .common import LayerNorm2d
class MaskDecoder(nn.Module):
def __init__(
self,
*,
transformer_dim: int,
transformer: nn.Module,
num_multimask_outputs: int = 3,
activation: Type[nn.Module] = nn.GELU,
iou_head_depth: int = 3,
iou_head_hidden_dim: int = 256,
use_high_res_features: bool = False,
iou_prediction_use_sigmoid=False,
dynamic_multimask_via_stability=False,
dynamic_multimask_stability_delta=0.05,
dynamic_multimask_stability_thresh=0.98,
pred_obj_scores: bool = False,
pred_obj_scores_mlp: bool = False,
use_multimask_token_for_obj_ptr: bool = False,
) -> None:
"""
Predicts masks given an image and prompt embeddings, using a
transformer architecture.
Arguments:
transformer_dim (int): the channel dimension of the transformer
transformer (nn.Module): the transformer used to predict masks
num_multimask_outputs (int): the number of masks to predict
when disambiguating masks
activation (nn.Module): the type of activation to use when
upscaling masks
iou_head_depth (int): the depth of the MLP used to predict
mask quality
iou_head_hidden_dim (int): the hidden dimension of the MLP
used to predict mask quality
"""
super().__init__()
self.transformer_dim = transformer_dim
self.transformer = transformer
self.num_multimask_outputs = num_multimask_outputs
self.iou_token = nn.Embedding(1, transformer_dim)
self.num_mask_tokens = num_multimask_outputs + 1
self.mask_tokens = nn.Embedding(self.num_mask_tokens, transformer_dim)
self.pred_obj_scores = pred_obj_scores
if self.pred_obj_scores:
self.obj_score_token = nn.Embedding(1, transformer_dim)
self.use_multimask_token_for_obj_ptr = use_multimask_token_for_obj_ptr
self.output_upscaling = nn.Sequential(
nn.ConvTranspose2d(
transformer_dim, transformer_dim // 4, kernel_size=2, stride=2
),
LayerNorm2d(transformer_dim // 4),
activation(),
nn.ConvTranspose2d(
transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2
),
activation(),
)
self.use_high_res_features = use_high_res_features
if use_high_res_features:
self.conv_s0 = nn.Conv2d(
transformer_dim, transformer_dim // 8, kernel_size=1, stride=1
)
self.conv_s1 = nn.Conv2d(
transformer_dim, transformer_dim // 4, kernel_size=1, stride=1
)
self.output_hypernetworks_mlps = nn.ModuleList(
[
MLP(transformer_dim, transformer_dim, transformer_dim // 8, 3)
for i in range(self.num_mask_tokens)
]
)
self.iou_prediction_head = MLP(
transformer_dim,
iou_head_hidden_dim,
self.num_mask_tokens,
iou_head_depth,
sigmoid_output=iou_prediction_use_sigmoid,
)
if self.pred_obj_scores:
self.pred_obj_score_head = nn.Linear(transformer_dim, 1)
if pred_obj_scores_mlp:
self.pred_obj_score_head = MLP(transformer_dim, transformer_dim, 1, 3)
# When outputting a single mask, optionally we can dynamically fall back to the best
# multimask output token if the single mask output token gives low stability scores.
self.dynamic_multimask_via_stability = dynamic_multimask_via_stability
self.dynamic_multimask_stability_delta = dynamic_multimask_stability_delta
self.dynamic_multimask_stability_thresh = dynamic_multimask_stability_thresh
def forward(
self,
image_embeddings: torch.Tensor,
image_pe: torch.Tensor,
sparse_prompt_embeddings: torch.Tensor,
dense_prompt_embeddings: torch.Tensor,
multimask_output: bool,
repeat_image: bool,
high_res_features: Optional[List[torch.Tensor]] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Predict masks given image and prompt embeddings.
Arguments:
image_embeddings (torch.Tensor): the embeddings from the image encoder
image_pe (torch.Tensor): positional encoding with the shape of image_embeddings
sparse_prompt_embeddings (torch.Tensor): the embeddings of the points and boxes
dense_prompt_embeddings (torch.Tensor): the embeddings of the mask inputs
multimask_output (bool): Whether to return multiple masks or a single
mask.
Returns:
torch.Tensor: batched predicted masks
torch.Tensor: batched predictions of mask quality
torch.Tensor: batched SAM token for mask output
"""
masks, iou_pred, mask_tokens_out, object_score_logits = self.predict_masks(
image_embeddings=image_embeddings,
image_pe=image_pe,
sparse_prompt_embeddings=sparse_prompt_embeddings,
dense_prompt_embeddings=dense_prompt_embeddings,
repeat_image=repeat_image,
high_res_features=high_res_features,
)
# Select the correct mask or masks for output
if multimask_output:
masks = masks[:, 1:, :, :]
iou_pred = iou_pred[:, 1:]
elif self.dynamic_multimask_via_stability and not self.training:
masks, iou_pred = self._dynamic_multimask_via_stability(masks, iou_pred)
else:
masks = masks[:, 0:1, :, :]
iou_pred = iou_pred[:, 0:1]
if multimask_output and self.use_multimask_token_for_obj_ptr:
sam_tokens_out = mask_tokens_out[:, 1:] # [b, 3, c] shape
else:
# Take the mask output token. Here we *always* use the token for single mask output.
# At test time, even if we track after 1-click (and using multimask_output=True),
# we still take the single mask token here. The rationale is that we always track
# after multiple clicks during training, so the past tokens seen during training
# are always the single mask token (and we'll let it be the object-memory token).
sam_tokens_out = mask_tokens_out[:, 0:1] # [b, 1, c] shape
# Prepare output
return masks, iou_pred, sam_tokens_out, object_score_logits
def predict_masks(
self,
image_embeddings: torch.Tensor,
image_pe: torch.Tensor,
sparse_prompt_embeddings: torch.Tensor,
dense_prompt_embeddings: torch.Tensor,
repeat_image: bool,
high_res_features: Optional[List[torch.Tensor]] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Predicts masks. See 'forward' for more details."""
# Concatenate output tokens
s = 0
if self.pred_obj_scores:
output_tokens = torch.cat(
[
self.obj_score_token.weight,
self.iou_token.weight,
self.mask_tokens.weight,
],
dim=0,
)
s = 1
else:
output_tokens = torch.cat(
[self.iou_token.weight, self.mask_tokens.weight], dim=0
)
output_tokens = output_tokens.unsqueeze(0).expand(
sparse_prompt_embeddings.size(0), -1, -1
)
tokens = torch.cat((output_tokens, sparse_prompt_embeddings), dim=1)
# Expand per-image data in batch direction to be per-mask
if repeat_image:
src = torch.repeat_interleave(image_embeddings, tokens.shape[0], dim=0)
else:
assert image_embeddings.shape[0] == tokens.shape[0]
src = image_embeddings
src = src + dense_prompt_embeddings
assert (
image_pe.size(0) == 1
), "image_pe should have size 1 in batch dim (from `get_dense_pe()`)"
pos_src = torch.repeat_interleave(image_pe, tokens.shape[0], dim=0)
b, c, h, w = src.shape
# Run the transformer
hs, src = self.transformer(src, pos_src, tokens)
iou_token_out = hs[:, s, :]
mask_tokens_out = hs[:, s + 1 : (s + 1 + self.num_mask_tokens), :]
# Upscale mask embeddings and predict masks using the mask tokens
src = src.transpose(1, 2).view(b, c, h, w)
if not self.use_high_res_features:
upscaled_embedding = self.output_upscaling(src)
else:
dc1, ln1, act1, dc2, act2 = self.output_upscaling
feat_s0, feat_s1 = high_res_features
upscaled_embedding = act1(ln1(dc1(src) + feat_s1))
upscaled_embedding = act2(dc2(upscaled_embedding) + feat_s0)
hyper_in_list: List[torch.Tensor] = []
for i in range(self.num_mask_tokens):
hyper_in_list.append(
self.output_hypernetworks_mlps[i](mask_tokens_out[:, i, :])
)
hyper_in = torch.stack(hyper_in_list, dim=1)
b, c, h, w = upscaled_embedding.shape
masks = (hyper_in @ upscaled_embedding.view(b, c, h * w)).view(b, -1, h, w)
# Generate mask quality predictions
iou_pred = self.iou_prediction_head(iou_token_out)
if self.pred_obj_scores:
assert s == 1
object_score_logits = self.pred_obj_score_head(hs[:, 0, :])
else:
# Obj scores logits - default to 10.0, i.e. assuming the object is present, sigmoid(10)=1
object_score_logits = 10.0 * iou_pred.new_ones(iou_pred.shape[0], 1)
return masks, iou_pred, mask_tokens_out, object_score_logits
def _get_stability_scores(self, mask_logits):
"""
Compute stability scores of the mask logits based on the IoU between upper and
lower thresholds.
"""
mask_logits = mask_logits.flatten(-2)
stability_delta = self.dynamic_multimask_stability_delta
area_i = torch.sum(mask_logits > stability_delta, dim=-1).float()
area_u = torch.sum(mask_logits > -stability_delta, dim=-1).float()
stability_scores = torch.where(area_u > 0, area_i / area_u, 1.0)
return stability_scores
def _dynamic_multimask_via_stability(self, all_mask_logits, all_iou_scores):
"""
When outputting a single mask, if the stability score from the current single-mask
output (based on output token 0) falls below a threshold, we instead select from
multi-mask outputs (based on output token 1~3) the mask with the highest predicted
IoU score. This is intended to ensure a valid mask for both clicking and tracking.
"""
# The best mask from multimask output tokens (1~3)
multimask_logits = all_mask_logits[:, 1:, :, :]
multimask_iou_scores = all_iou_scores[:, 1:]
best_scores_inds = torch.argmax(multimask_iou_scores, dim=-1)
batch_inds = torch.arange(
multimask_iou_scores.size(0), device=all_iou_scores.device
)
best_multimask_logits = multimask_logits[batch_inds, best_scores_inds]
best_multimask_logits = best_multimask_logits.unsqueeze(1)
best_multimask_iou_scores = multimask_iou_scores[batch_inds, best_scores_inds]
best_multimask_iou_scores = best_multimask_iou_scores.unsqueeze(1)
# The mask from singlemask output token 0 and its stability score
singlemask_logits = all_mask_logits[:, 0:1, :, :]
singlemask_iou_scores = all_iou_scores[:, 0:1]
stability_scores = self._get_stability_scores(singlemask_logits)
is_stable = stability_scores >= self.dynamic_multimask_stability_thresh
# Dynamically fall back to best multimask output upon low stability scores.
mask_logits_out = torch.where(
is_stable[..., None, None].expand_as(singlemask_logits),
singlemask_logits,
best_multimask_logits,
)
iou_scores_out = torch.where(
is_stable.expand_as(singlemask_iou_scores),
singlemask_iou_scores,
best_multimask_iou_scores,
)
return mask_logits_out, iou_scores_out
# Lightly adapted from
# https://github.com/facebookresearch/MaskFormer/blob/main/mask_former/modeling/transformer/transformer_predictor.py # noqa
class MLP(nn.Module):
def __init__(
self,
input_dim: int,
hidden_dim: int,
output_dim: int,
num_layers: int,
sigmoid_output: bool = False,
) -> None:
super().__init__()
self.num_layers = num_layers
h = [hidden_dim] * (num_layers - 1)
self.layers = nn.ModuleList(
nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])
)
self.sigmoid_output = sigmoid_output
def forward(self, x):
for i, layer in enumerate(self.layers):
x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
if self.sigmoid_output:
x = F.sigmoid(x)
return x
sam3/sam/prompt_encoder.py
+243
View File
@@ -0,0 +1,243 @@
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
from typing import Any, Optional, Tuple, Type
import numpy as np
import torch
from torch import nn
from .common import LayerNorm2d
class PromptEncoder(nn.Module):
def __init__(
self,
embed_dim: int,
image_embedding_size: Tuple[int, int],
input_image_size: Tuple[int, int],
mask_in_chans: int,
activation: Type[nn.Module] = nn.GELU,
) -> None:
"""
Encodes prompts for input to SAM's mask decoder.
Arguments:
embed_dim (int): The prompts' embedding dimension
image_embedding_size (tuple(int, int)): The spatial size of the
image embedding, as (H, W).
input_image_size (int): The padded size of the image as input
to the image encoder, as (H, W).
mask_in_chans (int): The number of hidden channels used for
encoding input masks.
activation (nn.Module): The activation to use when encoding
input masks.
"""
super().__init__()
self.embed_dim = embed_dim
self.input_image_size = input_image_size
self.image_embedding_size = image_embedding_size
self.pe_layer = PositionEmbeddingRandom(embed_dim // 2)
self.num_point_embeddings: int = 4 # pos/neg point + 2 box corners
point_embeddings = [
nn.Embedding(1, embed_dim) for i in range(self.num_point_embeddings)
]
self.point_embeddings = nn.ModuleList(point_embeddings)
self.not_a_point_embed = nn.Embedding(1, embed_dim)
self.mask_input_size = (
4 * image_embedding_size[0],
4 * image_embedding_size[1],
)
self.mask_downscaling = nn.Sequential(
nn.Conv2d(1, mask_in_chans // 4, kernel_size=2, stride=2),
LayerNorm2d(mask_in_chans // 4),
activation(),
nn.Conv2d(mask_in_chans // 4, mask_in_chans, kernel_size=2, stride=2),
LayerNorm2d(mask_in_chans),
activation(),
nn.Conv2d(mask_in_chans, embed_dim, kernel_size=1),
)
self.no_mask_embed = nn.Embedding(1, embed_dim)
def get_dense_pe(self) -> torch.Tensor:
"""
Returns the positional encoding used to encode point prompts,
applied to a dense set of points the shape of the image encoding.
Returns:
torch.Tensor: Positional encoding with shape
1x(embed_dim)x(embedding_h)x(embedding_w)
"""
return self.pe_layer(self.image_embedding_size).unsqueeze(0)
def _embed_points(
self,
points: torch.Tensor,
labels: torch.Tensor,
pad: bool,
) -> torch.Tensor:
"""Embeds point prompts."""
points = points + 0.5 # Shift to center of pixel
if pad:
padding_point = torch.zeros((points.shape[0], 1, 2), device=points.device)
padding_label = -torch.ones((labels.shape[0], 1), device=labels.device)
points = torch.cat([points, padding_point], dim=1)
labels = torch.cat([labels, padding_label], dim=1)
point_embedding = self.pe_layer.forward_with_coords(
points, self.input_image_size
)
point_embedding = torch.where(
(labels == -1).unsqueeze(-1),
torch.zeros_like(point_embedding) + self.not_a_point_embed.weight,
point_embedding,
)
point_embedding = torch.where(
(labels == 0).unsqueeze(-1),
point_embedding + self.point_embeddings[0].weight,
point_embedding,
)
point_embedding = torch.where(
(labels == 1).unsqueeze(-1),
point_embedding + self.point_embeddings[1].weight,
point_embedding,
)
point_embedding = torch.where(
(labels == 2).unsqueeze(-1),
point_embedding + self.point_embeddings[2].weight,
point_embedding,
)
point_embedding = torch.where(
(labels == 3).unsqueeze(-1),
point_embedding + self.point_embeddings[3].weight,
point_embedding,
)
return point_embedding
def _embed_boxes(self, boxes: torch.Tensor) -> torch.Tensor:
"""Embeds box prompts."""
boxes = boxes + 0.5 # Shift to center of pixel
coords = boxes.reshape(-1, 2, 2)
corner_embedding = self.pe_layer.forward_with_coords(
coords, self.input_image_size
)
corner_embedding[:, 0, :] += self.point_embeddings[2].weight
corner_embedding[:, 1, :] += self.point_embeddings[3].weight
return corner_embedding
def _embed_masks(self, masks: torch.Tensor) -> torch.Tensor:
"""Embeds mask inputs."""
mask_embedding = self.mask_downscaling(masks)
return mask_embedding
def _get_batch_size(
self,
points: Optional[Tuple[torch.Tensor, torch.Tensor]],
boxes: Optional[torch.Tensor],
masks: Optional[torch.Tensor],
) -> int:
"""
Gets the batch size of the output given the batch size of the input prompts.
"""
if points is not None:
return points[0].shape[0]
elif boxes is not None:
return boxes.shape[0]
elif masks is not None:
return masks.shape[0]
else:
return 1
def _get_device(self) -> torch.device:
return self.point_embeddings[0].weight.device
def forward(
self,
points: Optional[Tuple[torch.Tensor, torch.Tensor]],
boxes: Optional[torch.Tensor],
masks: Optional[torch.Tensor],
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Embeds different types of prompts, returning both sparse and dense
embeddings.
Arguments:
points (tuple(torch.Tensor, torch.Tensor) or none): point coordinates
and labels to embed.
boxes (torch.Tensor or none): boxes to embed
masks (torch.Tensor or none): masks to embed
Returns:
torch.Tensor: sparse embeddings for the points and boxes, with shape
BxNx(embed_dim), where N is determined by the number of input points
and boxes.
torch.Tensor: dense embeddings for the masks, in the shape
Bx(embed_dim)x(embed_H)x(embed_W)
"""
bs = self._get_batch_size(points, boxes, masks)
sparse_embeddings = torch.empty(
(bs, 0, self.embed_dim), device=self._get_device()
)
if points is not None:
coords, labels = points
point_embeddings = self._embed_points(coords, labels, pad=(boxes is None))
sparse_embeddings = torch.cat([sparse_embeddings, point_embeddings], dim=1)
if boxes is not None:
box_embeddings = self._embed_boxes(boxes)
sparse_embeddings = torch.cat([sparse_embeddings, box_embeddings], dim=1)
if masks is not None:
dense_embeddings = self._embed_masks(masks)
else:
dense_embeddings = self.no_mask_embed.weight.reshape(1, -1, 1, 1).expand(
bs, -1, self.image_embedding_size[0], self.image_embedding_size[1]
)
return sparse_embeddings, dense_embeddings
class PositionEmbeddingRandom(nn.Module):
"""
Positional encoding using random spatial frequencies.
"""
def __init__(self, num_pos_feats: int = 64, scale: Optional[float] = None) -> None:
super().__init__()
if scale is None or scale <= 0.0:
scale = 1.0
self.register_buffer(
"positional_encoding_gaussian_matrix",
scale * torch.randn((2, num_pos_feats)),
)
def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor:
"""Positionally encode points that are normalized to [0,1]."""
# assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape
coords = 2 * coords - 1
coords = coords @ self.positional_encoding_gaussian_matrix
coords = 2 * np.pi * coords
# outputs d_1 x ... x d_n x C shape
return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1)
def forward(self, size: Tuple[int, int]) -> torch.Tensor:
"""Generate positional encoding for a grid of the specified size."""
h, w = size
device: Any = self.positional_encoding_gaussian_matrix.device
grid = torch.ones((h, w), device=device, dtype=torch.float32)
y_embed = grid.cumsum(dim=0) - 0.5
x_embed = grid.cumsum(dim=1) - 0.5
y_embed = y_embed / h
x_embed = x_embed / w
pe = self._pe_encoding(torch.stack([x_embed, y_embed], dim=-1))
return pe.permute(2, 0, 1) # C x H x W
def forward_with_coords(
self, coords_input: torch.Tensor, image_size: Tuple[int, int]
) -> torch.Tensor:
"""Positionally encode points that are not normalized to [0,1]."""
coords = coords_input.clone()
coords[:, :, 0] = coords[:, :, 0] / image_size[1]
coords[:, :, 1] = coords[:, :, 1] / image_size[0]
return self._pe_encoding(coords.to(torch.float)) # B x N x C
sam3/sam/rope.py
+161
View File
@@ -0,0 +1,161 @@
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
"""
Adapted from:
1. https://github.com/meta-llama/codellama/blob/main/llama/model.py
2. https://github.com/naver-ai/rope-vit
3. https://github.com/lucidrains/rotary-embedding-torch
"""
from typing import Optional
import torch
from einops import rearrange, repeat
from torch import broadcast_tensors, nn
def init_t_xy(end_x: int, end_y: int, scale: float = 1.0, offset: int = 0, device=None):
t = torch.arange(end_x * end_y, dtype=torch.float32, device=device)
t_x = (t % end_x).float()
t_y = torch.div(t, end_x, rounding_mode="floor").float()
return t_x * scale + offset, t_y * scale + offset
def compute_axial_cis(
dim: int,
end_x: int,
end_y: int,
theta: float = 10000.0,
scale_pos: float = 1.0,
offset: int = 0,
device=None,
):
freqs_x = 1.0 / (
theta ** (torch.arange(0, dim, 4, device=device)[: (dim // 4)].float() / dim)
)
freqs_y = 1.0 / (
theta ** (torch.arange(0, dim, 4, device=device)[: (dim // 4)].float() / dim)
)
t_x, t_y = init_t_xy(end_x, end_y, scale_pos, offset, device=device)
freqs_x = torch.outer(t_x, freqs_x)
freqs_y = torch.outer(t_y, freqs_y)
freqs_cis_x = torch.polar(torch.ones_like(freqs_x), freqs_x)
freqs_cis_y = torch.polar(torch.ones_like(freqs_y), freqs_y)
return torch.cat([freqs_cis_x, freqs_cis_y], dim=-1)
def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
ndim = x.ndim
assert 0 <= 1 < ndim
assert freqs_cis.shape == (x.shape[-2], x.shape[-1])
shape = [d if i >= ndim - 2 else 1 for i, d in enumerate(x.shape)]
return freqs_cis.view(*shape)
def apply_rotary_enc(
xq: torch.Tensor,
xk: torch.Tensor,
freqs_cis: torch.Tensor,
repeat_freqs_k: bool = False,
):
xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
xk_ = (
torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
if xk.shape[-2] != 0
else None
)
freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
if xk_ is None:
# no keys to rotate, due to dropout
return xq_out.type_as(xq).to(xq.device), xk
# repeat freqs along seq_len dim to match k seq_len
if repeat_freqs_k:
r = xk_.shape[-2] // xq_.shape[-2]
freqs_cis = freqs_cis.repeat(*([1] * (freqs_cis.ndim - 2)), r, 1)
xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
return xq_out.type_as(xq).to(xq.device), xk_out.type_as(xk).to(xk.device)
def complex_mult(xq_real, xq_imag, freqs_cis_real, freqs_cis_imag):
# Compute the real part of the product
real_part = xq_real * freqs_cis_real - xq_imag * freqs_cis_imag
# Compute the imaginary part of the product
imag_part = xq_real * freqs_cis_imag + xq_imag * freqs_cis_real
# Stack the real and imaginary parts along the last dimension
return torch.stack([real_part, imag_part], dim=-1)
def apply_rotary_enc_real(
xq: torch.Tensor,
xk: torch.Tensor,
freqs_cis_real: torch.Tensor,
freqs_cis_imag: torch.Tensor,
repeat_freqs_k: bool = False,
):
assert xk is not None
assert xk.shape[-2] != 0
xq_real = xq.float().reshape(*xq.shape[:-1], -1, 2)[..., 0]
xq_imag = xq.float().reshape(*xq.shape[:-1], -1, 2)[..., 1]
xk_real = xk.float().reshape(*xk.shape[:-1], -1, 2)[..., 0]
xk_imag = xk.float().reshape(*xk.shape[:-1], -1, 2)[..., 1]
freqs_cis_real = reshape_for_broadcast(freqs_cis_real, xq_real)
freqs_cis_imag = reshape_for_broadcast(freqs_cis_imag, xq_imag)
xq_out = complex_mult(xq_real, xq_imag, freqs_cis_real, freqs_cis_imag).flatten(3)
if repeat_freqs_k:
r = xk_real.shape[-2] // xq_real.shape[-2]
freqs_cis_real = freqs_cis_real.repeat(*([1] * (freqs_cis_real.ndim - 2)), r, 1)
freqs_cis_imag = freqs_cis_imag.repeat(*([1] * (freqs_cis_imag.ndim - 2)), r, 1)
xk_out = complex_mult(xk_real, xk_imag, freqs_cis_real, freqs_cis_imag).flatten(3)
# xq_out = torch.view_as_real(torch.complex(xq_real, xq_imag) * torch.complex(freqs_cis_real, freqs_cis_imag)).flatten(3)
# xk_out = torch.view_as_real(torch.compelx(xk_real, xk_imag) * torch.complex(freqs_cis_real, freqs_cis_imag)).flatten(3)
return xq_out.type_as(xq).to(xq.device), xk_out.type_as(xk).to(xk.device)
# rotary embedding helper functions
def broadcat(tensors, dim=-1):
broadcasted_tensors = broadcast_tensors(*tensors)
return torch.cat(broadcasted_tensors, dim=dim)
def rotate_half(x: torch.Tensor):
x = rearrange(x, "... (d r) -> ... d r", r=2)
x1, x2 = x.unbind(dim=-1)
x = torch.stack((-x2, x1), dim=-1)
return rearrange(x, "... d r -> ... (d r)")
class VisionRotaryEmbeddingVE(nn.Module):
def __init__(
self,
dim: int,
seq_len: int,
pt_seq_len: Optional[int] = None,
theta: float = 10000.0,
offset: int = 1, # specific to VE
):
super().__init__()
freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
scale = 1.0
if pt_seq_len is not None:
scale = pt_seq_len / seq_len
# offset of +1 following VE - even though for the
# attention op only differences matter
t = torch.arange(seq_len) * scale + offset
freqs = torch.einsum("..., f -> ... f", t, freqs)
freqs = repeat(freqs, "... n -> ... (n r)", r=2)
freqs = broadcat((freqs[None, :, :], freqs[:, None, :]), dim=-1)
freqs_cos = freqs.cos().view(-1, freqs.shape[-1])
freqs_sin = freqs.sin().view(-1, freqs.shape[-1])
self.register_buffer("freqs_cos", freqs_cos)
self.register_buffer("freqs_sin", freqs_sin)
def forward(self, t: torch.Tensor):
return t * self.freqs_cos + rotate_half(t) * self.freqs_sin
sam3/sam/transformer.py
+358
View File
@@ -0,0 +1,358 @@
# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
import math
from functools import partial
from typing import Tuple, Type
import torch
import torch.nn.functional as F
from sam3.sam.rope import apply_rotary_enc, apply_rotary_enc_real, compute_axial_cis
from torch import nn, Tensor
from .common import MLPBlock
class TwoWayTransformer(nn.Module):
def __init__(
self,
depth: int,
embedding_dim: int,
num_heads: int,
mlp_dim: int,
activation: Type[nn.Module] = nn.ReLU,
attention_downsample_rate: int = 2,
) -> None:
"""
A transformer decoder that attends to an input image using
queries whose positional embedding is supplied.
Args:
depth (int): number of layers in the transformer
embedding_dim (int): the channel dimension for the input embeddings
num_heads (int): the number of heads for multihead attention. Must
divide embedding_dim
mlp_dim (int): the channel dimension internal to the MLP block
activation (nn.Module): the activation to use in the MLP block
"""
super().__init__()
self.depth = depth
self.embedding_dim = embedding_dim
self.num_heads = num_heads
self.mlp_dim = mlp_dim
self.layers = nn.ModuleList()
for i in range(depth):
self.layers.append(
TwoWayAttentionBlock(
embedding_dim=embedding_dim,
num_heads=num_heads,
mlp_dim=mlp_dim,
activation=activation,
attention_downsample_rate=attention_downsample_rate,
skip_first_layer_pe=(i == 0),
)
)
self.final_attn_token_to_image = Attention(
embedding_dim, num_heads, downsample_rate=attention_downsample_rate
)
self.norm_final_attn = nn.LayerNorm(embedding_dim)
def forward(
self,
image_embedding: Tensor,
image_pe: Tensor,
point_embedding: Tensor,
) -> Tuple[Tensor, Tensor]:
"""
Args:
image_embedding (torch.Tensor): image to attend to. Should be shape
B x embedding_dim x h x w for any h and w.
image_pe (torch.Tensor): the positional encoding to add to the image. Must
have the same shape as image_embedding.
point_embedding (torch.Tensor): the embedding to add to the query points.
Must have shape B x N_points x embedding_dim for any N_points.
Returns:
torch.Tensor: the processed point_embedding
torch.Tensor: the processed image_embedding
"""
# BxCxHxW -> BxHWxC == B x N_image_tokens x C
bs, c, h, w = image_embedding.shape
image_embedding = image_embedding.flatten(2).permute(0, 2, 1)
image_pe = image_pe.flatten(2).permute(0, 2, 1)
# Prepare queries
queries = point_embedding
keys = image_embedding
# Apply transformer blocks and final layernorm
for layer in self.layers:
queries, keys = layer(
queries=queries,
keys=keys,
query_pe=point_embedding,
key_pe=image_pe,
)
# Apply the final attention layer from the points to the image
q = queries + point_embedding
k = keys + image_pe
attn_out = self.final_attn_token_to_image(q=q, k=k, v=keys)
queries = queries + attn_out
queries = self.norm_final_attn(queries)
return queries, keys
class TwoWayAttentionBlock(nn.Module):
def __init__(
self,
embedding_dim: int,
num_heads: int,
mlp_dim: int = 2048,
activation: Type[nn.Module] = nn.ReLU,
attention_downsample_rate: int = 2,
skip_first_layer_pe: bool = False,
) -> None:
"""
A transformer block with four layers: (1) self-attention of sparse
inputs, (2) cross attention of sparse inputs to dense inputs, (3) mlp
block on sparse inputs, and (4) cross attention of dense inputs to sparse
inputs.
Arguments:
embedding_dim (int): the channel dimension of the embeddings
num_heads (int): the number of heads in the attention layers
mlp_dim (int): the hidden dimension of the mlp block
activation (nn.Module): the activation of the mlp block
skip_first_layer_pe (bool): skip the PE on the first layer
"""
super().__init__()
self.self_attn = Attention(embedding_dim, num_heads)
self.norm1 = nn.LayerNorm(embedding_dim)
self.cross_attn_token_to_image = Attention(
embedding_dim, num_heads, downsample_rate=attention_downsample_rate
)
self.norm2 = nn.LayerNorm(embedding_dim)
self.mlp = MLPBlock(embedding_dim, mlp_dim, activation)
self.norm3 = nn.LayerNorm(embedding_dim)
self.norm4 = nn.LayerNorm(embedding_dim)
self.cross_attn_image_to_token = Attention(
embedding_dim, num_heads, downsample_rate=attention_downsample_rate
)
self.skip_first_layer_pe = skip_first_layer_pe
def forward(
self, queries: Tensor, keys: Tensor, query_pe: Tensor, key_pe: Tensor
) -> Tuple[Tensor, Tensor]:
# Self attention block
if self.skip_first_layer_pe:
queries = self.self_attn(q=queries, k=queries, v=queries)
else:
q = queries + query_pe
attn_out = self.self_attn(q=q, k=q, v=queries)
queries = queries + attn_out
queries = self.norm1(queries)
# Cross attention block, tokens attending to image embedding
q = queries + query_pe
k = keys + key_pe
attn_out = self.cross_attn_token_to_image(q=q, k=k, v=keys)
queries = queries + attn_out
queries = self.norm2(queries)
# MLP block
mlp_out = self.mlp(queries)
queries = queries + mlp_out
queries = self.norm3(queries)
# Cross attention block, image embedding attending to tokens
q = queries + query_pe
k = keys + key_pe
attn_out = self.cross_attn_image_to_token(q=k, k=q, v=queries)
keys = keys + attn_out
keys = self.norm4(keys)
return queries, keys
class Attention(nn.Module):
"""
An attention layer that allows for downscaling the size of the embedding
after projection to queries, keys, and values.
"""
def __init__(
self,
embedding_dim: int,
num_heads: int,
downsample_rate: int = 1,
dropout: float = 0.0,
kv_in_dim: int = None,
use_fa3: bool = False,
) -> None:
super().__init__()
self.embedding_dim = embedding_dim
self.kv_in_dim = kv_in_dim if kv_in_dim is not None else embedding_dim
self.internal_dim = embedding_dim // downsample_rate
self.num_heads = num_heads
self.use_fa3 = use_fa3
assert (
self.internal_dim % num_heads == 0
), "num_heads must divide embedding_dim."
self.q_proj = nn.Linear(embedding_dim, self.internal_dim)
self.k_proj = nn.Linear(self.kv_in_dim, self.internal_dim)
self.v_proj = nn.Linear(self.kv_in_dim, self.internal_dim)
self.out_proj = nn.Linear(self.internal_dim, embedding_dim)
self.dropout_p = dropout
def _separate_heads(self, x: Tensor, num_heads: int) -> Tensor:
b, n, c = x.shape
x = x.reshape(b, n, num_heads, c // num_heads)
return x.transpose(1, 2) # B x N_heads x N_tokens x C_per_head
def _recombine_heads(self, x: Tensor) -> Tensor:
b, n_heads, n_tokens, c_per_head = x.shape
x = x.transpose(1, 2)
return x.reshape(b, n_tokens, n_heads * c_per_head) # B x N_tokens x C
def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
# Input projections
q = self.q_proj(q)
k = self.k_proj(k)
v = self.v_proj(v)
# Separate into heads
q = self._separate_heads(q, self.num_heads)
k = self._separate_heads(k, self.num_heads)
v = self._separate_heads(v, self.num_heads)
dropout_p = self.dropout_p if self.training else 0.0
# Attention
# with torch.backends.cuda.sdp_kernel(
# enable_flash=USE_FLASH_ATTN,
# # if Flash attention kernel is off, then math kernel needs to be enabled
# enable_math=(OLD_GPU and dropout_p > 0.0) or MATH_KERNEL_ON,
# enable_mem_efficient=OLD_GPU,
# ):
# Let's trust the dispatcher....
if self.use_fa3:
from sam3.perflib.fa3 import flash_attn_func
assert dropout_p == 0.0
out = flash_attn_func(
q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
).transpose(1, 2)
else:
torch.backends.cuda.enable_flash_sdp(True)
torch.backends.cuda.enable_math_sdp(True)
torch.backends.cuda.enable_mem_efficient_sdp(True)
out = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p)
out = self._recombine_heads(out)
out = self.out_proj(out)
return out
class RoPEAttention(Attention):
"""Attention with rotary position encoding."""
def __init__(
self,
*args,
rope_theta=10000.0,
# whether to repeat q rope to match k length
# this is needed for cross-attention to memories
rope_k_repeat=False,
feat_sizes=(64, 64), # [w, h] for stride 16 feats at 1024 resolution
use_rope_real=False,
**kwargs,
):
super().__init__(*args, **kwargs)
self.use_rope_real = use_rope_real
self.compute_cis = partial(
compute_axial_cis, dim=self.internal_dim // self.num_heads, theta=rope_theta
)
device = torch.device("cuda") if torch.cuda.is_available() else None
self.freqs_cis = self.compute_cis(
end_x=feat_sizes[0], end_y=feat_sizes[1], device=device
)
if self.use_rope_real:
self.freqs_cis_real = self.freqs_cis.real
self.freqs_cis_imag = self.freqs_cis.imag
self.rope_k_repeat = rope_k_repeat
def forward(
self, q: Tensor, k: Tensor, v: Tensor, num_k_exclude_rope: int = 0
) -> Tensor:
# Input projections
q = self.q_proj(q)
k = self.k_proj(k)
v = self.v_proj(v)
# Separate into heads
q = self._separate_heads(q, self.num_heads)
k = self._separate_heads(k, self.num_heads)
v = self._separate_heads(v, self.num_heads)
# Apply rotary position encoding
w = h = math.sqrt(q.shape[-2])
if self.freqs_cis.shape[0] != q.shape[-2]:
self.freqs_cis = self.compute_cis(end_x=w, end_y=h, device=q.device)
self.freqs_cis_real = self.freqs_cis.real
self.freqs_cis_imag = self.freqs_cis.imag
if q.shape[-2] != k.shape[-2]:
assert self.rope_k_repeat
num_k_rope = k.size(-2) - num_k_exclude_rope
if self.use_rope_real:
q, k[:, :, :num_k_rope] = apply_rotary_enc_real(
q,
k[:, :, :num_k_rope],
freqs_cis_real=self.freqs_cis_real,
freqs_cis_imag=self.freqs_cis_imag,
repeat_freqs_k=self.rope_k_repeat,
)
else:
q, k[:, :, :num_k_rope] = apply_rotary_enc(
q,
k[:, :, :num_k_rope],
self.freqs_cis,
repeat_freqs_k=self.rope_k_repeat,
)
dropout_p = self.dropout_p if self.training else 0.0
# Attention
# with torch.backends.cuda.sdp_kernel(
# enable_flash=USE_FLASH_ATTN,
# # if Flash attention kernel is off, then math kernel needs to be enabled
# enable_math=(OLD_GPU and dropout_p > 0.0) or MATH_KERNEL_ON,
# enable_mem_efficient=OLD_GPU,
# ):
# Let's trust the dispatcher....
if self.use_fa3:
from sam3.perflib.fa3 import flash_attn_func
assert dropout_p == 0.0
out = flash_attn_func(
q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
).transpose(1, 2)
else:
torch.backends.cuda.enable_flash_sdp(True)
torch.backends.cuda.enable_math_sdp(True)
torch.backends.cuda.enable_mem_efficient_sdp(True)
out = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p)
out = self._recombine_heads(out)
out = self.out_proj(out)
return out