Use torch to generate pos embeddings

Author: yizhuoz004

tripy/examples/segment-anything-model-v2/sam2/modeling/backbones/image_encoder.py

@@ -120,7 +120,8 @@ def __init__(
         self.position_encoding = []
         position_encoding_shapes = [[256, 256], [128, 128], [64, 64], [32, 32]]
         for s in position_encoding_shapes:
-            self.position_encoding.append(position_encoding.generate_static_embedding([1, 256] + s, dtype=dtype))
+            embed = position_encoding.generate_static_embedding([1, 256] + s, dtype=dtype)
+            self.position_encoding.append(tp.Tensor(embed))
 
     def __call__(self, xs: List[tp.Tensor]):
 

tripy/examples/segment-anything-model-v2/sam2/modeling/memory_encoder.py

@@ -179,6 +179,4 @@ def forward(
         x = self.fuser(x)
         x = self.out_proj(x)
 
-        pos = tp.cast(self.position_encoding(x), x.dtype)
-
-        return x, pos
+        return x

tripy/examples/segment-anything-model-v2/sam2/modeling/position_encoding.py

@@ -80,7 +80,7 @@ def forward(self, x: tp.Tensor):
         pos = tp.permute(pos, (0, 3, 1, 2))
         return pos
 
-    def generate_static_embedding(self, inp_shape, dtype):
+    def generate_static_embedding(self, inp_shape, dtype=None):
         import torch
 
         B, _, H, W = inp_shape
@@ -100,7 +100,12 @@ def generate_static_embedding(self, inp_shape, dtype):
         pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
         pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
         pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
-        return tp.Tensor(pos.to(getattr(torch, dtype)).contiguous())
+        if dtype is not None:
+            pos = pos.to(getattr(torch, dtype))
+        return pos.contiguous()
+
+    def generate_pos_embedding_torch(self, x):
+        return self.generate_static_embedding(x.shape).to(x.dtype).contiguous()
 
 
 class PositionEmbeddingRandom(tp.Module):

tripy/examples/segment-anything-model-v2/sam2/modeling/sam2_base.py

@@ -139,6 +139,7 @@ def __init__(
 
         # Part 3: memory encoder for the previous frame's outputs
         self.memory_encoder = memory_encoder
+        self.position_encoder = self.memory_encoder.position_encoding
         self.mem_dim = self.hidden_dim
         if hasattr(self.memory_encoder, "out_proj") and hasattr(self.memory_encoder.out_proj, "weight"):
             # if there is compression of memories along channel dim
@@ -720,11 +721,11 @@ def _encode_new_memory(
         if self.sigmoid_bias_for_mem_enc != 0.0:
             mask_for_mem = mask_for_mem + self.sigmoid_bias_for_mem_enc
 
-        maskmem_features, maskmem_pos_enc = self.memory_encoder(
+        maskmem_features = self.memory_encoder(
             tp.Tensor(pix_feat.float().contiguous()), tp.Tensor(mask_for_mem.contiguous())
         )  # sigmoid already applied
         maskmem_features = torch.from_dlpack(maskmem_features)
-        maskmem_pos_enc = [torch.from_dlpack(maskmem_pos_enc)]
+        maskmem_pos_enc = [self.position_encoder.generate_pos_embedding_torch(maskmem_features)]
 
         return maskmem_features, maskmem_pos_enc