:(

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import requests
 import json
+import time
 
 url = "https://nexra.aryahcr.cc/api/image/complements"
 headers = {
@@ -14,23 +15,53 @@ def generate_image(prompt, 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."
+            response_data = response.json()
+            image_id = response_data.get("id")
+            if not image_id:
+                return "Error: No image ID returned in the response."
+
+            while True:
+                status_response = requests.get(f"{url}/{image_id}")
+                if status_response.status_code == 200:
+                    status_data = status_response.json()
+                    status = status_data.get("status")
+                    if status == "completed":
+                        images = status_data.get("images")
+                        if images and isinstance(images, list):
+                            image_data = images[0]
+                            if image_data.startswith("data:image"):
+                                base64_data = image_data.split(",", 1)[1]
+                                return base64_data
+                            else:
+                                return "Error: Unexpected image data format."
+                        else:
+                            return "Error: No images found in the response."
+                    elif status == "error":
+                        return "Error: Image generation failed."
+                    elif status == "not_found":
+                        return "Error: Image ID not found."
+                    elif status == "pending":
+                        time.sleep(1)
+                    else:
+                        return f"Error: Unexpected status '{status}'."
+                else:
+                    return f"Error: Status check failed with code {status_response.status_code}."
         else:
-            return f"Error: {response.status_code}"
+            return f"Error: Initial request failed with code {response.status_code}."
     except json.JSONDecodeError:
         return "Error: Invalid JSON response."
     except Exception as e:
-        return str(e)
+        return f"Exception occurred: {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 = gr.Interface(
+    fn=generate_image,
+    inputs=[
+        gr.Textbox(label="Enter prompt"),
+        gr.Radio(["dalle2"], label="Select Model")
+    ],
+    outputs=gr.Image(type="pil"),
+    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()