Update app.py

app.py

@@ -2,6 +2,7 @@ import os
 import random
 import requests
 import gradio as gr
+from PIL import Image
 from io import BytesIO
 
 def generate_image(prompt, width, height, seed, randomize):
@@ -17,7 +18,11 @@ def generate_image(prompt, width, height, seed, randomize):
         .replace("[seed]", str(seed))
     )
     response = requests.get(url)
-    return BytesIO(response.content)
+    if response.status_code == 200:
+        image = Image.open(BytesIO(response.content))
+        return image
+    else:
+        return "Failed to retrieve image."
 
 title = "Unlimited FLUX-Pro"
 description = "Here you can use the FLUX-Pro as much as you want without limits"