import streamlit as st
import requests
from PIL import Image, UnidentifiedImageError
from transformers import AutoProcessor, Blip2ForConditionalGeneration
import torch
import io


st.title("Image Captioning with Fine-Tuned BLiPv2 Model")

uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    try:
        image = Image.open(uploaded_file)
        st.image(image, caption="Uploaded Image", use_column_width=True)
        
        try:
            files = {"file": uploaded_file.getvalue()}
            print("Sending API request")
            response = requests.post("http://0.0.0.0:8502/generate-caption/", files=files)
            caption = response.json().get("caption")
        except requests.exceptions.RequestException as e:
            st.error(f"An error occurred while making the API request: {str(e)}")
            caption = "Error generating caption"
        except ValueError as e:
            st.error(f"An error occurred while parsing the API response: {str(e)}")
            caption = "Error generating caption"

        st.write("Generated Caption:")
        st.write(f"**{caption}**")
        
    except UnidentifiedImageError:
        st.error("The uploaded file is not a valid image. Please upload a JPG, JPEG, or PNG file.")
    except Exception as e:
        st.error(f"An unexpected error occurred: {str(e)}")
else:
    st.write("Please upload an image to generate a caption.")