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

def generate_image(api_key, positive_prompt, negative_prompt, controlnet_image):
    # API request
    resp = requests.post(
        "https://model-5wov7jkq.api.baseten.co/production/predict",
        headers={"Authorization": f"Api-Key {api_key}"},
        json={
            'workflow_values': {
                'negative_prompt': negative_prompt,
                'positive_prompt': positive_prompt,
                'controlnet_image': controlnet_image
            }
        }
    )
    
    # Check if the request was successful
    if resp.status_code != 200:
        return "Error: Failed to get a valid response from the API."
    
    # Parse the response
    result = resp.json()["result"][0]
    
    # Decode base64 image data
    image_data = base64.b64decode(result["data"])
    
    # Create a PIL Image object
    image = Image.open(io.BytesIO(image_data))
    
    return image

# Create Gradio interface
iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="API Key", value="5ShnuaQF.N8bwpmZgbaAH2l21v813tUPA7RlX2rx9"),
        gr.Textbox(label="Positive Prompt", value="An igloo on a snowy day, 4k, hd"),
        gr.Textbox(label="Negative Prompt", value="blurry, text, low quality"),
        gr.Textbox(label="ControlNet Image URL", value="https://storage.googleapis.com/logos-bucket-01/baseten_logo.png")
    ],
    outputs=gr.Image(type="pil"),
    title="Image Generation UI",
    description="Generate an image using the provided API"
)

# Launch the interface
iface.launch()