Update vision_tower_builder.py

vision_tower_builder.py

@@ -24,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:
-    print("You need to install flash_attn")
-from timm.layers import drop_path, to_2tuple, trunc_normal_
+    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:
+    from timm.models.layers import drop_path, trunc_normal_, to_2tuple
 
 
 
@@ -67,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
@@ -613,7 +627,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}")
+    raise ValueError(f"Unknown vision tower: {vision_tower}")