Update vision_tower_builder.py

vision_tower_builder.py

@@ -24,8 +24,10 @@ import torch.utils.checkpoint as checkpoint
 from functools import partial
 try:
     from flash_attn import flash_attn_qkvpacked_func
+    use_flash_attn = True
 except:
-    print("You need to install flash_attn")
+    use_flash_attn = False
+    print("You need to install flash_attn to be faster!")
 from timm.layers import drop_path, to_2tuple, trunc_normal_
 
 
@@ -67,6 +69,12 @@ 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'
+            
         super().__init__()
         self.num_heads = num_heads
         head_dim = dim // num_heads
@@ -613,7 +621,6 @@ def build_vision_tower(vision_tower_cfg, **kwargs):
     if "umt-hd" in vision_tower:
         return UMTVisionTower(vision_tower, vision_tower_cfg=vision_tower_cfg, image_size=448, **kwargs)
     elif "umt" in vision_tower:
-        raise NotImplementedError
         return UMTVisionTower(vision_tower, vision_tower_cfg=vision_tower_cfg, **kwargs)
 
     raise ValueError(f"Unknown vision tower: {vision_tower}")