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

# Load the model and processor from Hugging Face Hub
model_name = "underthehoodst/cartoon-captioning"
processor = BlipProcessor.from_pretrained(model_name)
model = BlipForConditionalGeneration.from_pretrained(model_name)

device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

st.title("Cartoon Caption Generator")

uploaded_file = st.file_uploader("Upload a cartoon image", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
    image = Image.open(uploaded_file).convert("RGB")
    st.image(image, caption="Uploaded Cartoon", use_container_width=True)  # Updated parameter
    st.write("Generating caption...")
    
    inputs = processor(images=image, return_tensors="pt").to(device)
    output_ids = model.generate(**inputs)
    caption = processor.decode(output_ids[0], skip_special_tokens=True)
    caption = ". ".join(sentence.strip().capitalize() for sentence in caption.split(". "))
    
    st.write(f"Generated Caption: {caption}")