import gradio as gr
import torch
import os
import requests
import io
from PIL import Image
API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
headers = {"Authorization": "Bearer " + os.getenv("HF_TOKEN")}
def infer(prompt):
    def query(payload):
    	response = requests.post(API_URL, headers=headers, json=payload)
    	return response.content
    image_bytes = query({
    	"inputs": prompt,
    })

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(" # Text-to-Image Gradio Template")

        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )
            run_button = gr.Button("Run", scale=0, variant="primary")
            result = gr.Image(label="Result", show_label=False)



        # Run inference when run_button is clicked
        run_button.click(
            infer,
            inputs=[
                prompt
            ],
            outputs=[result],
        )

if __name__ == "__main__":
    demo.launch()