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

token = os.environ['AUTH_TOKEN']
runpod_id = os.environ['RUNPOD_ID']

url = 'https://{runpod_id}-8000.proxy.runpod.net/generate-image/'

def generate_image_from_text(prompt, class_guidance, token):
    headers = {
        'Authorization': f'Bearer {token}'
    }
    
    data = {
        "prompt": prompt,
        "class_guidance": class_guidance,
        "seed": 11,
        "num_imgs": 4,
        "img_size": 32
    }
    
    response = requests.post(url, json=data, headers=headers)
    
    if response.status_code == 200:
        image = Image.open(BytesIO(response.content))
    else:
        print("Failed to fetch image:", response.status_code, response.text)

    return image

# Define the Gradio interface
iface = gr.Interface(
    fn=generate_image_from_text,  # The function to generate the image
    inputs=["text", "slider"],
    outputs="image",
    title="Text-to-Image Generator",
    description="Enter a text prompt to generate an image."
)

# Launch the app
iface.launch()