Update app.py

app.py

@@ -1,12 +1,29 @@
 import gradio as gr
-from PIL import Image, ImageDraw
+import random
+from PIL import Image
+import io
 
 def generate_image(prompt):
-    # Create a simple image for testing
-    img = Image.new('RGB', (100, 100), color='red')
-    draw = ImageDraw.Draw(img)
-    draw.text((10, 10), prompt, fill='white')
-    return img
+    try:
+        # Generate a random seed for each request
+        random_seed = random.randint(0, 10000)
+        
+        # Generate image with the model
+        image = model.generate(prompt, seed=random_seed)
+        
+        # Check if the image is a PIL Image, if not try to convert
+        if isinstance(image, bytes):
+            image = Image.open(io.BytesIO(image))
+        elif not isinstance(image, Image.Image):
+            # If image is not in PIL format, convert or handle it
+            return "Error: Image not in expected format"
+        
+        return image
+    except Exception as e:
+        return f"Error: {e}"
+
+# Load the model
+model = gr.load("models/black-forest-labs/FLUX.1-schnell")
 
 # Create Gradio interface
 iface = gr.Interface(fn=generate_image, inputs="text", outputs="image")