fix err

app.py

@@ -4,39 +4,64 @@ import json
 import time
 
 url = "https://nexra.aryahcr.cc/api/image/complements"
-headers = {"Content-Type": "application/json"}
-
+headers = {
+    "Content-Type": "application/json"
+}
 
 def generate_image(prompt, model):
-    data = {"prompt": prompt, "model": model}
-    response = requests.post(url, headers=headers, data=json.dumps(data))
-
-    if response.status_code == 200:
-        response_data = response.json()
-        image_id = response_data.get("id")
+    data = {
+        "prompt": prompt,
+        "model": model,
+    }
+    try:
+        response = requests.post(url, headers=headers, data=json.dumps(data))
+        
+        if response.status_code == 200:
+            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")
+            while True:
+                status_response = requests.get(f"{url}/{image_id}")
                 
-                if status == "completed":
-                    images = status_data.get("images")
-                    if images:
-                        return images[0]
-                    return no_nsfw_image_path
-                elif status in ["error", "not_found"]:
-                    return no_nsfw_image_path
-                time.sleep(1)
-    return no_nsfw_image_path
+                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_url = images[0]
+                            return image_url
+                        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: Initial request failed with code {response.status_code}."
+    except json.JSONDecodeError:
+        return "Error: Invalid JSON response."
+    except Exception as 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")],
+    inputs=[
+        gr.Textbox(label="Enter prompt"),
+        gr.Radio(["dalle2"], label="Select Model")
+    ],
     outputs=gr.Image(type="filepath"),
     title="DALLE2 Generation",
     description="DALLE-2 Generation from Nextra API, NOTE: DALL-E 3 is not supported yet."
 )
-
-iface.launch(share=True)
+iface.launch()