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