import gradio as gr
import requests
import json

url = "https://nexra.aryahcr.cc/api/image/complements"
headers = {
    "Content-Type": "application/json"
}

def generate_image(prompt, model):
    data = {
        "prompt": prompt,
        "model": model,
    }
    try:
        response = requests.post(url, headers=headers, data=json.dumps(data))
        raw_response = response.text.strip()

        if response.status_code == 200:
            cleaned_response = raw_response.lstrip('_')
            response_data = json.loads(cleaned_response)
            
            if 'images' in response_data and isinstance(response_data['images'], list):
                image_url = response_data['images'][0]
                return image_url
            else:
                return "Error: No image data in the response."
        else:
            return f"Error: {response.status_code}"
    except json.JSONDecodeError:
        return "Error: Invalid JSON response."
    except Exception as e:
        return str(e)


iface = gr.Interface(fn=generate_image, inputs=[gr.Textbox(label="Enter prompt"), gr.Radio(["dalle2"], label="Select Model")], outputs="image", title="DALLE2 Generation", description="Disclaimer: This uses the Nexra API for image generation. I cannot guarantee rate limits, if you do not recieve an image please try again. I will change to use a different API soon. DALL-E 3 is not supported yet.")
iface.launch()