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

API_URL = "https://api-inference.huggingface.co/models/runwayml/stable-diffusion-v1-5"
TOKEN = os.getenv("HUGGING_FACE_TOKEN")

def query(payload):
    headers = {"Authorization": f"Bearer {TOKEN}"}
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.content

def text_to_image(text, temperature=0.7, top_p=0.9):
    payload = {
        "inputs": text,
        "parameters": {
            "temperature": temperature,
            "top_p": top_p
        }
    }
    image_bytes = query(payload)
    image = Image.open(io.BytesIO(image_bytes))
    return image

# Define Gradio interface
text_to_image_interface = gr.Interface(
    fn=text_to_image,
    inputs=[
        gr.inputs.Textbox(lines=2, label="Enter text"),
        gr.inputs.Slider(minimum=0.1, maximum=2, default=0.7, label="Temperature", step=0.1),
        gr.inputs.Slider(minimum=0.1, maximum=1, default=0.9, label="Top P", step=0.1)
    ],
    outputs=gr.outputs.Image(label="Generated Image"),
    title="Text to Image Generator",
    description="Enter text to generate an image.",
    examples=[["Astronaut riding a horse"]]
)

# Launch the interface
text_to_image_interface.launch()