import gradio as gr
import requests
import io
import os
from PIL import Image

# Get the API URL and headers from environment variables
API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
headers = {"Authorization": f"Bearer {os.getenv('HF_API_TOKEN')}"}

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

# Define the function to generate the image
def generate_image(input_text):
    image_bytes = query({"inputs": input_text})
    image = Image.open(io.BytesIO(image_bytes))
    return image

# Create Gradio interface
iface = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="Input Text", placeholder="Enter your description here..."),
    outputs=gr.Image(label="Generated Image"),
    title="LodhranGPT Image Generator",
    description="Enter a description to generate an image."
)

# Launch the Gradio app
iface.launch()