import streamlit as st
import requests
import time
import os
import random
from PIL import Image
from io import BytesIO
import io

HF_TOKEN = os.getenv("HF_TOKEN")

def query_model(text_input):
    API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
    payload = {
        "inputs": text_input,
        "num_inference_steps": 20,
        "guidance_scale": 4.5,
        "seed": random.randint(1, 9999999),
        "width": 1024,
        "height": 1024,
        "negative_prompt": "blurry, ugly, deformed, bad anatomy"
    }
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.content

def sdxl():
    st.title("Image Creator")
    text_input = st.text_input("Enter your prompt:", "Astronaut riding a horse")

    generated_image = None

    if st.button("Create"):   
        image_bytes = query_model(text_input)
        generated_image = Image.open(io.BytesIO(image_bytes))

    if generated_image is not None:
        st.image(generated_image, use_column_width=True)