Update app.py

app.py

@@ -1,28 +1,59 @@
 import gradio as gr
 from transformers import BlipProcessor, BlipForConditionalGeneration
 from PIL import Image
+import base64
+import io
 
 # Load processor and model
 processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
 model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
 
-# Captioning function
+# Captioning function for direct image input
 def caption_image(image):
     inputs = processor(images=image, return_tensors="pt")
     out = model.generate(**inputs)
     caption = processor.decode(out[0], skip_special_tokens=True)
     return caption
 
-# Interface
-iface = gr.Interface(
-    fn=caption_image,
-    inputs=gr.Image(type="pil"),
-    outputs="text",
-    title="Explain this Image",
-    flagging_mode="never",
-    # Add the api_name parameter
-    api_name="predict"
-)
+# API endpoint function that can handle base64 images
+def api_caption_image(base64_img):
+    try:
+        # Remove the data URL prefix if present
+        if "," in base64_img:
+            base64_img = base64_img.split(",")[1]
+        
+        # Decode base64 to image
+        image_bytes = base64.b64decode(base64_img)
+        image = Image.open(io.BytesIO(image_bytes))
+        
+        # Process with model
+        inputs = processor(images=image, return_tensors="pt")
+        out = model.generate(**inputs)
+        caption = processor.decode(out[0], skip_special_tokens=True)
+        return caption
+    except Exception as e:
+        return f"Error processing image: {str(e)}"
 
-# Queue for API & launch
-iface.queue(api_open=True).launch(share=True)
+# Create Blocks for more flexibility
+with gr.Blocks() as demo:
+    with gr.Tab("Image Captioning"):
+        gr.Interface(
+            fn=caption_image,
+            inputs=gr.Image(type="pil"),
+            outputs="text",
+            title="Explain this Image",
+            flagging_mode="never",
+        )
+    
+    # Define the API endpoint explicitly
+    gr.Interface(
+        fn=api_caption_image,
+        inputs=gr.Textbox(),  # For base64 input
+        outputs="text",
+        title="API Endpoint",
+        flagging_mode="never",
+        api_name="predict"  # This is the API endpoint name
+    )
+
+# Launch with queue and API open
+demo.queue(api_open=True).launch(share=True)