import streamlit as st
from freeGPT import Client
from PIL import Image
from io import BytesIO
import os
import requests

st.sidebar.title("Stable Diffusion SDXL-Turbo", help="Made with API")
st.sidebar.subheader("By KVI Kontent")
st.sidebar.write("Choose model and enter prompt")
api_key = os.environ['api_key']
API_URL_DALLE = "https://api-inference.huggingface.co/models/openskyml/dalle-3-xl"
headers = {"Authorization": f"Bearer {api_key}"}

model = st.sidebar.selectbox("Choose Model", ("prodia", "pollinations", "Dall-e 3"))

prompt = st.sidebar.text_input("Prompt", "")

def Dalle_query(payload):
    response = requests.post(API_URL_DALLE, headers=headers, json=payload)
    return response.content

if model == "prodia" or model == "pollinations":
    try:
        resp = Client.create_generation(model, prompt)
        image = Image.open(BytesIO(resp))
        st.image(image, caption="Generated Image")
    except Exception as e:
        st.error(str(e))
elif model == "Dall-e 3":
    image_bytes = Dalle_query({
    	"inputs": prompt
    })
    image = Image.open(BytesIO(image_bytes))
    st.image(image, caption="Generated Image")