Update app.py

app.py

@@ -1,65 +1,29 @@
-from dotenv import load_dotenv, find_dotenv
+# app.py
+
 import os
 import io
 from PIL import Image
-import requests
 import warnings
 import gradio as gr
+from transformers import pipeline
 
-# Suppress specific warnings
+# Suppress warnings
 warnings.filterwarnings("ignore", message=".*Using the model-agnostic default max_length.*")
 
-# Load environment variables from .env file
-load_dotenv(find_dotenv())
-hf_api_key = os.getenv('HF_API_KEY')
-endpoint_url = os.getenv('HF_API_ITT_BASE')
-
-# Validate environment variables
-if not hf_api_key:
-    raise ValueError("HF_API_KEY is not set in the .env file.")
-if not endpoint_url:
-    raise ValueError("HF_API_ITT_BASE is not set in the .env file.")
-
+# Load BLIP image captioning pipeline
+captioner = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
 
-def generate_caption(image):
-    """
-    Sends an image to the Hugging Face Inference Endpoint for caption generation.
-    Sends raw image bytes (not base64).
-    :param image: An image in PIL format.
-    :return: Generated caption or error message.
-    """
+# Function to generate caption using the pipeline
+def generate_caption(image: Image.Image):
     try:
-        headers = {"Authorization": f"Bearer {hf_api_key}"}
-
-        # Convert the image to RGB and save as JPEG into buffer
-        buffered = io.BytesIO()
+        # Convert image to RGB just in case
         image = image.convert("RGB")
-        image.save(buffered, format="JPEG")
-        buffered.seek(0)
-
-        # Send raw image bytes to the endpoint
-        response = requests.post(endpoint_url, headers=headers, data=buffered.getvalue())
-
-        # Try to parse JSON response safely
-        try:
-            response_data = response.json()
-        except ValueError:
-            return f"Invalid response (not JSON): {response.text}"
-
-        if response.status_code == 200:
-            if isinstance(response_data, list):
-                return response_data[0].get("generated_text", "No caption generated.")
-            elif isinstance(response_data, dict):
-                return response_data.get("generated_text", "No caption generated.")
-            else:
-                return f"Unexpected response format: {response_data}"
-        else:
-            return f"Error {response.status_code}: {response.text}"
-
+        # Generate caption
+        caption = captioner(image)[0]["generated_text"]
+        return caption
     except Exception as e:
         return f"An error occurred: {str(e)}"
 
-
 # Predefined sample images
 def get_sample_images():
     """
@@ -75,9 +39,10 @@ def get_sample_images():
     except FileNotFoundError:
         return []
 
+# Load sample images
+sample_images = get_sample_images()
 
 # Gradio interface
-sample_images = get_sample_images()  # Load predefined sample images
 demo = gr.Interface(
     fn=generate_caption,
     inputs=gr.Image(type="pil", label="Upload Image"),
@@ -86,12 +51,11 @@ demo = gr.Interface(
     title="Image Captioning App",
     description=(
         "Upload an image or use one of the predefined samples to generate a caption. "
-        "This app uses a Hugging Face Inference Endpoint for the Salesforce/blip-image-captioning-large model."
+        "This app uses `Salesforce/blip-image-captioning-base` locally via Hugging Face Transformers."
     ),
+    flagging_mode="never"
 )
 
 if __name__ == "__main__":
-    # Launch the Gradio demo
     demo.launch()
 
-