import streamlit as st 
import requests
import io
from PIL import Image
import os 

hf_token = os.environ.get("hf_token")
API_URL_KVI = "https://api-inference.huggingface.co/models/Kvikontent/kviimager2.0"
API_URL_MJ = "https://api-inference.huggingface.co/models/Kvikontent/midjourney-v6"
API_URL_DALLE = "https://api-inference.huggingface.co/models/ehristoforu/dalle-3-xl"
headers = {"Authorization": f"Bearer {hf_token}"}

def query(payload):
	response = requests.post(API_URL, headers=headers, json=payload)
	return response.content

st.title("Text To Image models")
st.write("Choose model and enter a prompt")
model = st.selectbox(
    "Choose model",
    ("KVIImager2.0", "Midjourney V6", "Dalle 3")
)
prompt = st.text_input("Enter prompt")

if prompt:
    if model == "KVIImager2.0":
        API_URL = API_URL_KVI
    elif model == "Midjourney V6":
        API_URL = API_URL_MJ
    elif model == "Dalle 3":
        API_URL = API_URL_DALLE
    image_bytes = query({
    	"inputs": prompt,
    })
    st.image(image_bytes, caption="Generated Image")
    st.info("Image generated successfully!")