verzv

app.py

@@ -20,7 +20,7 @@ from PIL import Image
 import torch
 import subprocess
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
-
+import time
 
 
 
@@ -119,7 +119,7 @@ def search(query: str, ds, images, k):
     device = "cuda:0" if torch.cuda.is_available() else "cpu"
     if device != model.device:
         model.to(device)
-        
+    start = time.time()    
     qs = []
     with torch.no_grad():
         batch_query = process_queries(processor, [query], mock_image)
@@ -135,6 +135,7 @@ def search(query: str, ds, images, k):
     results = []
     for idx in top_k_indices:
         results.append((images[idx])) #, f"Page {idx}"
+    print(f"Temps: {time.time()- start} s")
     del model
     del processor
     print("done")