import einops
import streamlit as st
from PIL import Image
from transformers import AutoModelForCausalLM, AutoTokenizer

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)

    model_id = "vikhyatk/moondream2"
    revision = "2024-05-08"
    model = AutoModelForCausalLM.from_pretrained(
        model_id, trust_remote_code=True, revision=revision
    )
    tokenizer = AutoTokenizer.from_pretrained(model_id, revision=revision)

    enc_image = model.encode_image(image)
    print(model.answer_question(enc_image, "Describe this image in a few sentences.", tokenizer))