Update image_pipeline.py

image_pipeline.py

@@ -1,16 +1,23 @@
-from transformers import pipeline
-
-def load_image_model():
-    """
-    Loads HuggingFaceTB/SmolVLM-500M-Instruct or another image-to-text model.
-    """
-    return pipeline("image-to-text", model="HuggingFaceTB/SmolVLM-500M-Instruct")
-
-def analyze_image(image_file, image_model):
-    """
-    Pass an image file to the image model pipeline and return the text/caption.
-    """
-    result = image_model(image_file)
-    if isinstance(result, list) and len(result) > 0:
-        return result[0].get("generated_text", "No caption generated.")
-    return "Unable to process image."
+import streamlit as st
+from transformers import pipeline
+from config import IMAGE_MODEL_NAME
+
+@st.cache_resource
+def load_image_model():
+    """
+    Loads an image captioning model recognized by the HF pipeline.
+    Example: "nlpconnect/vit-gpt2-image-captioning" or "Salesforce/blip-image-captioning-base".
+    """
+    return pipeline("image-to-text", model=IMAGE_MODEL_NAME)
+
+def analyze_image(image_file, image_model):
+    """
+    Pass an uploaded image to the loaded pipeline for caption generation.
+    """
+    try:
+        result = image_model(image_file)
+        if isinstance(result, list) and len(result) > 0:
+            return result[0].get("generated_text", "No caption.")
+        return "No output from the model."
+    except Exception as e:
+        return f"Error analyzing image: {str(e)}"