Add implement of tiny and update env

Signed-off-by: GitHub <noreply@github.com>

extra/grpc/sam/sam.py

@@ -17,7 +17,7 @@ import grpc
 import torch
 from functools import partial
 from segment_anything_hq import SamAutomaticMaskGenerator
-from segment_anything_hq.modeling import ImageEncoderViT, MaskDecoderHQ, PromptEncoder, Sam, TwoWayTransformer
+from segment_anything_hq.modeling import ImageEncoderViT, MaskDecoderHQ, PromptEncoder, Sam, TwoWayTransformer, TinyViT
 import matplotlib.pyplot as plt
 import numpy as np
 
@@ -54,8 +54,8 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
             model_name = request.model_name
             if model_name not in SamModelType.__dict__.keys():
                 raise Exception(f"Model name {model_name} not found in {SamModelType.__dict__.keys()}")
+            
             model_path = request.model_path
-            # check the model_path is valid
             if not os.path.exists(model_path):
                 raise Exception(f"Model path {model_path} does not exist")
 
@@ -69,8 +69,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
                 case SamModelType.vit_b:
                     sam = _build_sam_vit_b(checkpoint=model_path)
                 case SamModelType.vit_tiny:
-                    # TODO: Implement this
-                    pass
+                    sam = _build_sam_vit_tiny(checkpoint=model_path)
                 case _:
                     raise Exception(f"Model name {model_name} not found in {SamModelType.__dict__.keys()}")
             # TODO No sure if this is the right way to do it
@@ -163,6 +162,57 @@ def _build_sam_vit_l(checkpoint=None):
 def _build_sam_vit_b(checkpoint=None):
     return _constrcut_sam(encoder_embed_dim=768,encoder_depth=12,encoder_num_heads=12,encoder_global_attn_indexes=[2,5,8,11],checkpoint=checkpoint)
 
+def _build_sam_vit_tiny(checkpoint=None):
+    image_embedding_size = IMAGE_SIZE // VIT_PATCH_SIZE
+    mobile_sam = Sam(
+            image_encoder=TinyViT(img_size=1024, in_chans=3, num_classes=1000,
+                embed_dims=[64, 128, 160, 320],
+                depths=[2, 2, 6, 2],
+                num_heads=[2, 4, 5, 10],
+                window_sizes=[7, 7, 14, 7],
+                mlp_ratio=4.,
+                drop_rate=0.,
+                drop_path_rate=0.0,
+                use_checkpoint=False,
+                mbconv_expand_ratio=4.0,
+                local_conv_size=3,
+                layer_lr_decay=0.8
+            ),
+            prompt_encoder=PromptEncoder(
+            embed_dim=PROMT_EMBED_DIM,
+            image_embedding_size=(image_embedding_size, image_embedding_size),
+            input_image_size=(IMAGE_SIZE, IMAGE_SIZE),
+            mask_in_chans=16,
+            ),
+            mask_decoder=MaskDecoderHQ(
+                    num_multimask_outputs=3,
+                    transformer=TwoWayTransformer(
+                    depth=2,
+                    embedding_dim=PROMT_EMBED_DIM,
+                    mlp_dim=2048,
+                    num_heads=8,
+                ),
+                transformer_dim=PROMT_EMBED_DIM,
+                iou_head_depth=3,
+                iou_head_hidden_dim=256,
+                vit_dim=160,
+            ),
+            pixel_mean=[123.675, 116.28, 103.53],
+            pixel_std=[58.395, 57.12, 57.375],
+        )
+
+    mobile_sam.eval()
+    if checkpoint is not None:
+        with open(checkpoint, "rb") as f:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            state_dict = torch.load(f, map_location=device)
+        info = mobile_sam.load_state_dict(state_dict, strict=False)
+        print(info)
+    for n, p in mobile_sam.named_parameters():
+        if 'hf_token' not in n and 'hf_mlp' not in n and 'compress_vit_feat' not in n and 'embedding_encoder' not in n and 'embedding_maskfeature' not in n:
+            p.requires_grad = False
+    return mobile_sam
+
 def masks_to_image(anns, request):
     if len(anns)==0:
         return