fix: throw warnings if xformers or flash-attn cant be used

modeling_clip.py

@@ -5,6 +5,8 @@
 # and adjusted for Jina CLIP
 
 import base64
+import importlib.util
+import warnings
 from functools import partial
 from io import BytesIO
 from typing import List, Optional, Tuple, Union
@@ -117,6 +119,61 @@ def _build_vision_tower(config: JinaCLIPVisionConfig) -> EVAVisionTransformer:
     )
 
 
+def _resolve_attention_libs(config: JinaCLIPConfig):
+    use_text_flash_attn = (
+        config.use_text_flash_attn
+        if config.use_text_flash_attn is not None
+        else config.text_config.hf_model_config_kwargs.get('use_flash_attn', True)
+    )
+    use_vision_xformers = (
+        config.use_vision_xformers
+        if config.use_vision_xformers is not None
+        else config.vision_config.x_attention
+    )
+
+    def _resolve_use_text_flash_attn() -> bool:
+        if use_text_flash_attn:
+            if not torch.cuda.is_available():
+                warnings.warn('Flash attention requires CUDA, disabling')
+                return False
+            if not importlib.util.find_spec('flash_attn') is None:
+                warnings.warn(
+                    'Flash attention is not installed. Check '
+                    'https://github.com/Dao-AILab/flash-attention?'
+                    'tab=readme-ov-file#installation-and-features '
+                    'for installation instructions, disabling'
+                )
+                return False
+            return True
+        return False
+
+    def _resolve_use_vision_xformers() -> bool:
+        if use_vision_xformers:
+            if not torch.cuda.is_available():
+                warnings.warn('xFormers requires CUDA, disabling')
+                return False
+            if not importlib.util.find_spec('xformers') is None:
+                warnings.warn(
+                    'xFormers is not installed. Check '
+                    'https://github.com/facebookresearch/xformers?'
+                    'tab=readme-ov-file#installing-xformers for installation '
+                    'instructions, disabling'
+                )
+                return False
+            return True
+        return False
+
+    _use_text_flash_attn = _resolve_use_text_flash_attn()
+    _use_vision_xformers = _resolve_use_vision_xformers()
+
+    config.use_text_flash_attn = _use_text_flash_attn
+    config.use_vision_xformers = _use_vision_xformers
+    config.text_config.hf_model_config_kwargs['use_flash_attn'] = _use_text_flash_attn
+    config.vision_config.x_attention = _use_vision_xformers
+
+    return config
+
+
 class JinaCLIPPreTrainedModel(PreTrainedModel):
     """
     An abstract class to handle weights initialization and a simple interface for
@@ -218,16 +275,10 @@ class JinaCLIPModel(JinaCLIPPreTrainedModel):
                 f'JinaCLIPVisionConfig but is of type {type(config.vision_config)}.'
             )
 
+        config = _resolve_attention_libs(config)
         text_config = config.text_config
         vision_config = config.vision_config
 
-        if config.use_text_flash_attn is not None:
-            text_config.hf_model_config_kwargs['use_flash_attn'] = (
-                config.use_text_flash_attn
-            )
-        if config.use_vision_xformers is not None:
-            vision_config.x_attention = config.use_vision_xformers
-
         self.add_projections = config.add_projections
         self.projection_dim = config.projection_dim
         self.text_embed_dim = text_config.embed_dim