import streamlit as st
from transformers import BlipProcessor, BlipForConditionalGeneration
from PIL import Image

# 1. Load the Model (Cached so it doesn't reload every time)
@st.cache_resource
def load_model():
    processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
    model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
    return processor, model

processor, model = load_model()

# 2. App UI Layout
st.title("📷 AI Image Caption Generator")
st.markdown("Upload an image, and this AI will describe what it sees.")

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

if uploaded_file is not None:
    # Display the image
    image = Image.open(uploaded_file).convert('RGB')
    st.image(image, caption='Uploaded Image', use_column_width=True)
    
    st.write("### 🤖 Generating Caption...")
    
    # 4. Generate Caption
    # Prepare image for the model
    inputs = processor(image, return_tensors="pt")
    
    # Generate output
    out = model.generate(**inputs)
    
    # Decode the output to text
    caption = processor.decode(out[0], skip_special_tokens=True)
    
    # Display result
    st.success(f"**Caption:** {caption.capitalize()}")