app.py

import streamlit as st
from transformers import BlipProcessor, BlipForConditionalGeneration
import torch
from PIL import Image
from backend import base_model, test_single_image, tokenizer

# Load the smart brain (model) and its helper (processor) once
@st.cache_resource()
def load_model():
    processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
    model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    return processor, model

#processor, model = load_model()

# Title for your website
st.title("Funny Image Caption Maker")

# Let the user upload a picture
uploaded_file = st.file_uploader("Upload a picture!", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
    # Open the picture
    image = Image.open(uploaded_file).convert("RGB")
    
    # Show the picture on the website
    st.image(image, caption="Your Uploaded Picture")

# If they upload something, do this
if st.button("Generate Caption"):
    with st.spinner():
            # Make a caption for the picture
            cap = test_single_image(model=base_model, image_path=uploaded_file,tokenizer=tokenizer)

    # Show the caption
    #st.write("Here’s your caption: ", caption)
    
    st.text_area("Here’s your caption: ", cap)