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

# Define a function to decode the base64 image and save it as a file
def save_image(base64_str, image_format='png'):
    image_data = base64.b64decode(base64_str)
    image = Image.open(io.BytesIO(image_data))
    # Generate a unique filename
    filename = f"output_image.{image_format}"
    image.save(filename)
    return filename

# Define the function that will call the API and handle the image
def query_api(negative_prompt, positive_prompt):
    # Prepare the headers and the JSON body of the request
    headers = {
        "Authorization": "Api-Key O7NlgQG7.9e3nxz7K4K3CakdBpSYkHNlGx5gOmaXb"
    }
    json_data = {
        'workflow_values': {
            'negative_prompt': negative_prompt,
            'positive_prompt': positive_prompt
        }
    }
    
    # Send the POST request
    response = requests.post(
        "https://model-5qe9kjp3.api.baseten.co/development/predict",
        headers=headers,
        json=json_data
    )
    
    # Process the response
    if response.status_code == 200:
        data = response.json()
        results = data.get('result', [])
        if results:
            # Assuming the first result is what we want to display
            base64_image = results[0].get('data')
            image_format = results[0].get('format', 'png')
            if base64_image and image_format:
                # Save the decoded image
                image_path = save_image(base64_image, image_format)
                return image_path
            else:
                return "No image data found in the API response."
        else:
            return "The API response did not contain any results."
    else:
        return "Failed to fetch data from the API."

# Define the Gradio interface
iface = gr.Interface(
    fn=query_api,
    inputs=[
        gr.Textbox(label="Negative Prompt", placeholder="Enter Negative Prompt Here"),
        gr.Textbox(label="Positive Prompt", placeholder="Enter Positive Prompt Here")
    ],
    outputs=gr.File(label="Download Image", type="filepath"),
    title="API Query Interface without ControlNet Image",
    description="Enter a negative prompt and a positive prompt to query the API and download the resulting image."
)

# Launch the Gradio app
iface.launch()