import requests
import gradio as gr
from PIL import Image
# from dotenv import load_dotenv
import os
import io
token = os.environ['HF_TOKEN']



# load_dotenv()
# token = os.environ.get("HF_TOKEN", "")
print(token)
API_URL = "https://api-inference.huggingface.co/models/playgroundai/playground-v2-1024px-aesthetic"
bearer_token = "Bearer " + str(token)
# token = 
headers = {"Authorization": bearer_token}

def query(payload):
    response = requests.post(API_URL, headers=headers, json=payload)
    image_bytes = response.content
    # image = Image.open(io.BytesIO(image_bytes))
    # return image_bytes
    image = Image.open(io.BytesIO(image_bytes))
    return image



with gr.Blocks() as demo:
    gr.Interface(fn=query, inputs="text", outputs="image")

demo.launch()