TAP v1.1 models release

app.py

@@ -32,7 +32,7 @@ def parse_args():
     """Parse arguments."""
     parser = argparse.ArgumentParser(description="Launch gradio application")
     parser.add_argument("--model-type", type=str, default="tap_vit_l")
-    parser.add_argument("--checkpoint", type=str, default="models/tap_vit_l_548184.pkl")
+    parser.add_argument("--checkpoint", type=str, default="models/tap_vit_l_v1_1.pkl")
     parser.add_argument("--concept", type=str, default="concepts/merged_2560.pkl")
     parser.add_argument("--device", nargs="+", type=int, default=[0], help="Index of devices")
     return parser.parse_args()

models/tap_vit_l_548184.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:e1d3a11c572af8cb6bce8016d3a6c6948bba4959ea43811f0e984b9eafeee413
-size 811637521

models/{tap_vit_l_03f8ec.pkl → tap_vit_l_v1_1.pkl}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d63a5aba993c34bf29c0466026136e18e25d2bd4ac9e51b8fc407b76c431707d
-size 811637521
+oid sha256:d7301f73267786bfafc7223931b3fd51baf111d1c24bcc1278ea6378139f1067
+size 811637487

tokenize_anything/layers/utils.py

@@ -26,7 +26,8 @@ def init_cross_conv(blocks):
         if isinstance(m, torch.nn.Conv2d):
             torch.nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
     for blk in blocks:
-        torch.nn.init.constant_(blk.norm3.weight, 0)
+        if hasattr(blk, "norm3") and hasattr(blk.norm3, "weight"):
+            torch.nn.init.zeros_(blk.norm3.weight)
 
 
 def set_dropout(module, dropout):

tokenize_anything/modeling/concept_projector.py

@@ -31,12 +31,12 @@ class ConceptProjector(nn.Module):
 
     def reset_weights(self, src_weights=None, tgt_weights=None):
         """Reset the normalized projection weights."""
-        if src_weights is not None:
+        if src_weights:
             with open(src_weights, "rb") as f:
                 self.src_weights, self.concepts = pickle.load(f)
                 self.src_weights = torch.from_numpy(self.src_weights)
                 self.concepts = np.array(self.concepts)
-        if tgt_weights is not None:
+        if tgt_weights:
             with open(tgt_weights, "rb") as f:
                 self.tgt_weights, self.concepts = pickle.load(f)
                 self.tgt_weights = torch.from_numpy(self.tgt_weights)

tokenize_anything/modeling/image_decoder.py

@@ -50,13 +50,12 @@ class Attention(nn.Module):
 
     def __init__(self, dim=256, num_heads=8, attn_ratio=1):
         super(Attention, self).__init__()
-        qkv_dim = int(dim * attn_ratio)
-        self.num_heads = num_heads
-        self.head_dim = qkv_dim // num_heads
-        self.q_proj = nn.Linear(dim, qkv_dim)
-        self.k_proj = nn.Linear(dim, qkv_dim)
-        self.v_proj = nn.Linear(dim, qkv_dim)
-        self.proj = nn.Linear(qkv_dim, dim)
+        self.num_heads = num_heads or dim // 64
+        self.head_dim = int(dim * attn_ratio) // self.num_heads
+        self.q_proj = nn.Linear(dim, self.num_heads * self.head_dim)
+        self.k_proj = nn.Linear(dim, self.num_heads * self.head_dim)
+        self.v_proj = nn.Linear(dim, self.num_heads * self.head_dim)
+        self.proj = nn.Linear(self.num_heads * self.head_dim, dim)
         self.scale = self.head_dim**-0.5
 
     def forward(self, q, k, v):
@@ -100,8 +99,7 @@ class Block(nn.Module):
         q, k = query + query_pos, key + key_pos
         query = self.norm2(self.dropout(self.cross_attn_token_to_image(q, k, key)).add_(query))
         query = self.norm3(self.dropout(self.mlp(query)).add_(query))
-        q = query + query_pos
-        key = self.norm4(self.cross_attn_image_to_token(k, q, query).add_(key))
+        key = self.norm4(self.cross_attn_image_to_token(k, query + query_pos, query).add_(key))
         return query, key
 
 
@@ -137,8 +135,7 @@ class Transformer(nn.Module):
         for blk in self.blocks:
             query, key = blk(query, key, query_pos, key_pos)
         q, k = query + query_pos, key + key_pos
-        query = self.dropout(self.final_attn_token_to_image(q, k, key)).add_(query)
-        query = self.norm(query)
+        query = self.norm(self.dropout(self.final_attn_token_to_image(q, k, key)).add_(query))
         return query, key
 
 
@@ -164,10 +161,10 @@ class ImageDecoder(nn.Module):
         super(ImageDecoder, self).__init__()
         self.embed_dim = embed_dim
         self.num_mask_tokens = num_mask_tokens
-        self.transformer = Transformer(embed_dim, num_heads=num_heads, depth=depth)
+        self.transformer = Transformer(embed_dim, num_heads, depth=depth)
         self.iou_token = nn.Embedding(1, embed_dim)
-        self.sem_tokens = nn.Embedding(self.num_mask_tokens, embed_dim)
-        self.mask_tokens = nn.Embedding(self.num_mask_tokens, embed_dim)
+        self.sem_tokens = nn.Embedding(num_mask_tokens, embed_dim)
+        self.mask_tokens = nn.Embedding(num_mask_tokens, embed_dim)
         self.output_conv = nn.Sequential(
             nn.ConvTranspose2d(embed_dim, embed_dim // 4, 2, 2),
             TransposedLayerNorm(embed_dim // 4),
@@ -178,8 +175,8 @@ class ImageDecoder(nn.Module):
         self.mask_pred = nn.ModuleList(
             Predictor(embed_dim, embed_dim // 8) for _ in range(num_mask_tokens)
         )
-        self.iou_pred = Predictor(embed_dim, self.num_mask_tokens)
-        self.sem_pred = Predictor(embed_dim, sem_embed_dim, 1024)
+        self.iou_pred = Predictor(embed_dim, num_mask_tokens)
+        self.sem_pred = Predictor(embed_dim, sem_embed_dim, sem_embed_dim)
 
     def get_outputs(self, inputs):
         img_embeds = inputs["img_embeds"]
@@ -201,18 +198,17 @@ class ImageDecoder(nn.Module):
         key = key.transpose(1, 2).view((-1, self.embed_dim) + img_embed_size)
         mask_embeds = self.output_conv(key).flatten(2)
         # Unpack query.
-        tokens = query[:, :num_tokens].unbind(dim=1)
-        iou_tokens = tokens[num_tokens - self.num_mask_tokens - 1]
-        mask_tokens = tokens[num_tokens - self.num_mask_tokens :]
-        sem_tokens = tokens[: self.num_mask_tokens]
+        sem_tokens = query[:, : self.num_mask_tokens]
+        sam_tokens = query[:, self.num_mask_tokens : num_tokens].unbind(1)
+        iou_tokens, mask_tokens = sam_tokens[0], sam_tokens[1:]
         # Predict.
         mask_pred = [f(x) for f, x in zip(self.mask_pred, mask_tokens)]
         mask_pred = torch.stack(mask_pred, dim=1) @ mask_embeds
         mask_pred_size = list(4 * embed_size for embed_size in img_embed_size)
         mask_pred = mask_pred.view([-1, self.num_mask_tokens] + mask_pred_size)
         outputs = {"iou_pred": self.iou_pred(iou_tokens), "mask_pred": mask_pred}
-        outputs["sem_tokens"] = torch.stack(sem_tokens, dim=1)
-        outputs["sem_embeds"] = self.sem_pred(outputs["sem_tokens"])
+        outputs["sem_tokens"] = sem_tokens.unsqueeze_(2)
+        outputs["sem_embeds"] = self.sem_pred(outputs["sem_tokens"].flatten(2))
         return outputs
 
     def forward(self, inputs):

tokenize_anything/modeling/prompt_encoder.py

@@ -24,21 +24,14 @@ class PromptEncoder(nn.Module):
 
     def __init__(self, embed_dim, image_size):
         super(PromptEncoder, self).__init__()
-        self.img_size = [image_size] * 2
         self.point_embed = nn.Embedding(5, embed_dim)  # [bg, fg, lt, rb, pad]
         self.corner_labels = torch.tensor([[2, 3]], dtype=torch.int64)
         self.register_buffer("coord_matrix", torch.randn((2, embed_dim // 2)))
-        self.img_pos = None
+        self.img_pos, self.img_size = None, [image_size] * 2
 
-    def to_tensor(self, input):
-        """Convert input to tensor."""
-        if input is None:
-            return input
-        if not isinstance(input, torch.Tensor):
-            input = torch.from_numpy(input)
-        if input.device != self.coord_matrix.device:
-            input = input.to(device=self.coord_matrix.device)
-        return input
+    def as_tensor(self, input):
+        """Convert input into a tensor."""
+        return torch.as_tensor(input, device=self.coord_matrix.device)
 
     def to_points(self, points=None, boxes=None):
         """Convert points or boxes to point prompts."""
@@ -48,14 +41,14 @@ class PromptEncoder(nn.Module):
             else:
                 coords, labels = points[:, :, :2], points[:, :, 2]
             coords = coords.__add__(0.5).__itruediv__(self.img_size[::-1])
-            coords = self.to_tensor(coords.clip(0, 1).astype("float32"))
-            labels = self.to_tensor(labels.astype("int64"))
+            coords = self.as_tensor(coords.clip(0, 1).astype("float32"))
+            labels = self.as_tensor(labels.astype("int64"))
             return coords, labels
         if boxes is not None:
             coords = boxes.reshape((-1, 2, 2))
             coords = coords.__add__(0.5).__itruediv__(self.img_size[::-1])
-            coords = self.to_tensor(coords.clip(0, 1).astype("float32"))
-            labels = self.to_tensor(self.corner_labels)
+            coords = self.as_tensor(coords.clip(0, 1).astype("float32"))
+            labels = self.as_tensor(self.corner_labels)
             return coords, labels
         return None
 
@@ -79,7 +72,7 @@ class PromptEncoder(nn.Module):
         grid = torch.ones(*grid_size, dtype=torch.float32)
         y = grid.cumsum(dim=0).sub_(0.5).div_(grid_size[0])
         x = grid.cumsum(dim=1).sub_(0.5).div_(grid_size[1])
-        coords = self.to_tensor(torch.stack([x, y], dim=-1))
+        coords = self.as_tensor(torch.stack([x, y], dim=-1))
         return self.encode_coords(coords)
 
     def forward(self, inputs):

tokenize_anything/models/easy_build.py

@@ -18,7 +18,6 @@
 from functools import partial
 import pickle
 
-import numpy as np
 import torch
 
 from tokenize_anything.modeling import ConceptProjector
@@ -45,9 +44,9 @@ def load_weights(module, weights_file, strict=True):
         with open(weights_file, "rb") as f:
             state_dict = pickle.load(f)
             for k, v in state_dict.items():
-                state_dict[k] = torch.from_numpy(v) if isinstance(v, np.ndarray) else v
+                state_dict[k] = torch.as_tensor(v)
     else:
-        state_dict = torch.load(weights_file)
+        state_dict = torch.load(weights_file, map_location="cpu")
     module.load_state_dict(state_dict, strict=strict)
 
 
@@ -68,21 +67,21 @@ def vit_encoder(depth, embed_dim, num_heads, out_dim, image_size):
 def image_tokenizer(image_encoder, checkpoint=None, device=0, dtype="float16", **kwargs):
     """Build an image tokenizer."""
     image_size = kwargs.get("image_size", 1024)
-    prompt_embed_dim = kwargs.get("prompt_embed_dim", 256)
+    image_embed_dim = kwargs.get("image_embed_dim", 256)
     sem_embed_dim = kwargs.get("sem_embed_dim", 1024)
     text_embed_dim = kwargs.get("text_embed_dim", 512)
     text_decoder_depth = kwargs.get("text_decoder_depth", 12)
     text_seq_len = kwargs.get("text_seq_len", 40)
     text_tokenizer = TextTokenizer()
     model = ImageTokenizer(
-        image_encoder=image_encoder(out_dim=prompt_embed_dim, image_size=image_size),
-        prompt_encoder=PromptEncoder(embed_dim=prompt_embed_dim, image_size=image_size),
+        image_encoder=image_encoder(out_dim=image_embed_dim, image_size=image_size),
+        prompt_encoder=PromptEncoder(embed_dim=image_embed_dim, image_size=image_size),
         image_decoder=ImageDecoder(
+            embed_dim=image_embed_dim,
+            num_heads=image_embed_dim // 32,
+            sem_embed_dim=sem_embed_dim,
             depth=2,
-            embed_dim=prompt_embed_dim,
-            num_heads=prompt_embed_dim // 32,
             num_mask_tokens=4,
-            sem_embed_dim=sem_embed_dim,
         ),
         text_tokenizer=text_tokenizer,
         concept_projector=ConceptProjector(),
@@ -90,10 +89,10 @@ def image_tokenizer(image_encoder, checkpoint=None, device=0, dtype="float16", *
             depth=text_decoder_depth,
             embed_dim=text_embed_dim,
             num_heads=text_embed_dim // 64,
-            mlp_ratio=4,
-            prompt_embed_dim=prompt_embed_dim,
+            prompt_embed_dim=image_embed_dim,
             max_seq_len=text_seq_len,
             vocab_size=text_tokenizer.n_words,
+            mlp_ratio=4,
         ),
     )
     load_weights(model, checkpoint)

tokenize_anything/version.py

@@ -1,3 +1,3 @@
-version = "0.1.0a0"
+version = "1.1.0a0"
 git_version = "None"
 __version__ = version