import io
import os
from tempfile import NamedTemporaryFile

import replicate
import streamlit as st
from PIL import Image
from dotenv import load_dotenv

load_dotenv()

os.environ["REPLICATE_API_TOKEN"] = os.getenv("REPLICATE_API_TOKEN")
MODEL = "daanelson/minigpt-4:b96a2f33cc8e4b0aa23eacfce731b9c41a7d9466d9ed4e167375587b54db9423"


st.title("Brighten Up My Day")

username = st.text_input("What's your name")
uploaded_file = st.file_uploader("Choose a file")
if uploaded_file is not None and username is not None:
    # To read file as bytes:
    image_data = uploaded_file.getvalue()
    image = Image.open(io.BytesIO(image_data))

    with NamedTemporaryFile(suffix=".jpg") as temp_file:
        image.save(temp_file.name)

        output = replicate.run(
            MODEL,
            input={
                "image": open(temp_file.name, "rb"),
                "temperature": 1.3,
                "prompt": f"This is {username}. Write something nice about {username} to brighten up their "
                          f"day. And give them compliments."
            }
        )

    st.divider()

    col1, col2 = st.columns(2)

    with col1:
        st.image(image)

    with col2:
        with st.container():
            st.write(output)