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

def call_api_and_generate_image(api_url, api_key, negative_prompt, positive_prompt, controlnet_image):
    if not api_key:
        return None, "Error: API key is required. Please provide a valid API key."

    if not api_url:
        return None, "Error: API URL is required. Please provide a valid API URL."

    resp = requests.post(
        api_url,
        headers={"Authorization": f"Api-Key {api_key}"},
        json={
            'workflow_values': {
                'negative_prompt': negative_prompt,
                'positive_prompt': positive_prompt,
                'controlnet_image': controlnet_image
            }
        }
    )

    if resp.status_code != 200:
        return None, f"Error: API call failed with status code {resp.status_code}, message: {resp.text}"

    result = resp.json()
    base64_data = result['result'][0]['data']
    image_data = base64.b64decode(base64_data)
    image = Image.open(io.BytesIO(image_data))
    return image, json.dumps(result, indent=2)

def generate_image(api_url, api_key, negative_prompt, positive_prompt, controlnet_image):
    try:
        return call_api_and_generate_image(api_url, api_key, negative_prompt, positive_prompt, controlnet_image)
    except Exception as e:
        return None, f"Error: {str(e)}"

# Create the Gradio interface
iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="API URL", placeholder="Enter the API URL here", value="https://your-baseten-url.baseten.co/development/predict"),
        gr.Textbox(label="API Key", placeholder="Enter your API key here"),
        gr.Textbox(label="Negative Prompt", value="blurry, text, low quality", placeholder="Enter negative prompt here..."),
        gr.Textbox(label="Positive Prompt", value="An igloo on a snowy day, 4k, hd", placeholder="Enter positive prompt here..."),
        gr.Textbox(label="ControlNet Image URL", value="https://storage.googleapis.com/logos-bucket-01/baseten_logo.png", placeholder="Enter ControlNet image URL here...")
    ],
    outputs=[
        gr.Image(type="pil", label="Generated Image"),
        gr.Textbox(label="API Response JSON", lines=10)
    ],
    title="Image Generation with API",
    description="Enter the API URL, your API key, and prompts to generate an image using the provided API."
)

# Launch the interface
iface.launch()