Update vision_tower_builder.py

vision_tower_builder.py

@@ -2,9 +2,6 @@ from typing import Optional, Tuple, Union, Dict
 from dataclasses import dataclass
 from functools import partial, reduce
 from PIL import Image
-import torch
-import torch.utils.checkpoint
-from torch import nn
 import os
 from transformers.image_processing_utils import BatchFeature, get_size_dict
 from transformers.image_transforms import (
@@ -27,9 +24,15 @@ import torch.utils.checkpoint as checkpoint
 from functools import partial
 try:
     from flash_attn import flash_attn_qkvpacked_func
+    use_flash_attn = True
+except:
+    use_flash_attn = False
+    print("You need to install flash_attn to be faster!")
+
+try:
+    from timm.layers import drop_path, to_2tuple, trunc_normal_
 except:
-    print("You need to install flash_attn")
-from timm.models.layers import drop_path, to_2tuple, trunc_normal_
+    from timm.models.layers import drop_path, trunc_normal_, to_2tuple
 
 
 
@@ -70,6 +73,14 @@ class Attention(nn.Module):
             self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.,
             proj_drop=0., attn_head_dim=None,
             attn_type='flash_v2'):
+
+        if use_flash_attn:
+            attn_type = attn_type
+        else:
+            attn_type = 'origin'
+        
+        print(attn_type)
+
         super().__init__()
         self.num_heads = num_heads
         head_dim = dim // num_heads
@@ -516,7 +527,7 @@ def build_vit(config, pt_type='origin'):
         drop_path_rate=0., 
         num_frames=config.num_frames,
         tubelet_size=1,
-        use_checkpoint=True,
+        use_checkpoint=False,
         checkpoint_num=24,
         return_index=config.return_idx,
         with_ln=True, # merge vision_layernorm in it
@@ -614,9 +625,8 @@ def build_vision_tower(vision_tower_cfg, **kwargs):
 
 
     if "umt-hd" in vision_tower:
-        raise NotImplementedError
         return UMTVisionTower(vision_tower, vision_tower_cfg=vision_tower_cfg, image_size=448, **kwargs)
     elif "umt" in vision_tower:
         return UMTVisionTower(vision_tower, vision_tower_cfg=vision_tower_cfg, **kwargs)
 
-    raise ValueError(f"Unknown vision tower: {vision_tower}")
+    raise ValueError(f"Unknown vision tower: {vision_tower}")