import streamlit as st
from PIL import Image
import torch

from transformers import AutoProcessor, AutoModelForVision2Seq

DEVICE = "cuda:0"

# Add a header
st.title("BEST Story Teller...Ever!")

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

if uploaded_file is not None:
    image = Image.open(uploaded_file)
    st.image(image, caption='Uploaded Image.', use_column_width=True)

    processor = AutoProcessor.from_pretrained("HuggingFaceM4/idefics2-8b-base")
    model = AutoModelForVision2Seq.from_pretrained(
        "HuggingFaceM4/idefics2-8b-base",
    ).to(DEVICE)

    # Create inputs
    prompts = [
        "Describe the image in a few sentences.",
    ]
    images = [image]
    inputs = processor(text=prompts, images=images, padding=True, return_tensors="pt")
    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}


    # Generate
    generated_ids = model.generate(**inputs, max_new_tokens=500)
    generated_texts = processor.batch_decode(generated_ids, skip_special_tokens=True)

    print(generated_texts)