Update inference.py

inference.py

@@ -6,6 +6,7 @@ import faiss
 import numpy as np
 import json
 import spacy
+import time
 
 # Load models and resources
 clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
@@ -112,10 +113,24 @@ def generate_rag_caption(image):
     Input: PIL Image or image path (str).
     Output: Caption (str).
     """
-    embedding = extract_image_features(image)
-    if embedding is None:
-        return "Failed to process image."
-    retrieved = retrieve_similar_captions(embedding, k=5)
-    if not retrieved:
-        return "No similar captions found."
-    return generate_caption_from_retrieved(retrieved)
+    try:
+        start = time.time()
+        embedding = extract_image_features(image)
+        print("⏱️ CLIP feature extraction:", round(time.time() - start, 2), "s")
+        if embedding is None:
+            return "Failed to process image."
+
+        start = time.time()
+        retrieved = retrieve_similar_captions(embedding, k=5)
+        print("⏱️ Caption retrieval:", round(time.time() - start, 2), "s")
+        if not retrieved:
+            return "No similar captions found."
+
+        start = time.time()
+        caption = generate_caption_from_retrieved(retrieved)
+        print("⏱️ Caption generation:", round(time.time() - start, 2), "s")
+
+        return caption
+    except Exception as e:
+        print(f"Error in RAG captioning: {e}")
+        return "Something went wrong during caption generation."