Update handler.py

handler.py

@@ -1,4 +1,5 @@
 from typing import Dict, List, Any
+from transformers.utils.import_utils import is_flash_attn_2_available
 from colpali_engine.models import ColQwen2, ColQwen2Processor
 import torch
 
@@ -8,7 +9,7 @@ class EndpointHandler():
             path,
             torch_dtype=torch.bfloat16,
             device_map="cuda:0",  # or "mps" if on Apple Silicon
-            # attn_implementation="flash_attention_2", # should work on A100
+            attn_implementation="flash_attention_2" if is_flash_attn_2_available() else None, # should work on A100
             ).eval()
         self.processor = ColQwen2Processor.from_pretrained(path)
 
@@ -24,7 +25,7 @@ class EndpointHandler():
         batch_images = self.processor.process_images([images]).to(self.model.device)
         # Forward pass
         with torch.no_grad():
-            image_embeddings = self.model(**batch_images)
+            image_embeddings = self.model(**batch_images).tolist()
 
-        return {"embeddings": image_embeddings.tolist()}
+        return {"embeddings": image_embeddings}