Update README.md

README.md

@@ -53,17 +53,17 @@ 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_name = "vidore/colqwen2-v1.0-merged"
-
 model = ColQwen2.from_pretrained(
-        model_name,
-        torch_dtype=torch.bfloat16,
-        device_map="cuda:0",  # or "mps" if on Apple Silicon
-    ).eval()
-processor = ColQwen2Processor.from_pretrained(model_name)
+    "vidore/colqwen2-v1.0-merged",
+    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-v1.0-merged")
 
 # Your inputs
 images = [