Update README.md

README.md

@@ -60,14 +60,16 @@ pip install git+https://github.com/illuin-tech/colpali
 ```python
 import torch
 from PIL import Image
+from transformers.utils.import_utils import is_flash_attn_2_available
 
 from colpali_engine.models import ColQwen2, ColQwen2Processor
 
 model = ColQwen2.from_pretrained(
-        "vidore/colqwen2-v0.1",
-        torch_dtype=torch.bfloat16,
-        device_map="cuda:0",  # or "mps" if on Apple Silicon
-    ).eval()
+    "vidore/colqwen2-v0.1",
+    torch_dtype=torch.bfloat16,
+    device_map="cuda:0",  # or "mps" if on Apple Silicon
+    attn_implementation="flash_attention_2" if is_flash_attn_2_available() else None,  # or "eager" if "mps"
+).eval()
 processor = ColQwen2Processor.from_pretrained("vidore/colqwen2-v0.1")
 
 # Your inputs