Update app.py

app.py

@@ -1,6 +1,9 @@
 import gradio as gr
 import replicate
 import os
+import requests
+from PIL import Image
+from io import BytesIO
 
 def generate_image(prompt, api_key):
     # Set the API key for the current session
@@ -18,8 +21,12 @@ def generate_image(prompt, api_key):
         input=inputs
     )
     
-    # Return the generated image URL
-    return output_url[0]
+    # Fetch the image from the URL
+    response = requests.get(output_url)
+    image = Image.open(BytesIO(response.content))
+    
+    # Return the image
+    return image
 
 # Create the Gradio interface
 iface = gr.Interface(
@@ -37,7 +44,7 @@ iface = gr.Interface(
             label="Replicate API Key"
         )
     ],
-    outputs=gr.Image(type="auto"),
+    outputs=gr.Image(type="pil"),
     title="FLUX 1.1 Pro Text-to-Image Generator",
     description="Generate images from text prompts using the FLUX 1.1 Pro model."
 )