Update app.py

app.py

@@ -1,83 +1,54 @@
 import gradio as gr
 import fal_client
-import time
 import requests
-from io import BytesIO
 from PIL import Image
-import base64
+from io import BytesIO
 
-def generate_image(api_key, prompt):
+def generate_image(api_key, prompt, image_size='landscape_4_3', num_images=1):
     try:
-        # Set the API key provided by the user
+        # Set the API key for the fal_client
         fal_client.api_key = api_key
 
-        # Default parameters
-        image_size = "landscape_4_3"
-        num_images = 1
-        enable_safety_checker = True
-        safety_tolerance = "2"
-        sync_mode = True  # Enable sync_mode to get the image directly
-
         handler = fal_client.submit(
             "fal-ai/flux-pro/v1.1",
             arguments={
                 "prompt": prompt,
                 "image_size": image_size,
                 "num_images": num_images,
-                "enable_safety_checker": enable_safety_checker,
-                "safety_tolerance": safety_tolerance,
-                "sync_mode": sync_mode,
             },
         )
-
-        # Since sync_mode is True, handler.get() will wait for the result
         result = handler.get()
-
-        # For debugging, print the result
-        print("API Response Result:", result)
-
-        if not result.get("images"):
-            return None, "No images were generated."
-
-        image_info = result["images"][0]
-
-        # Check if image data is in the 'data' field (base64 encoded)
-        if 'data' in image_info:
-            # Image data is base64 encoded
-            image_data = image_info['data']
-            image_bytes = base64.b64decode(image_data)
-            image = Image.open(BytesIO(image_bytes))
-            return image, None
-        elif 'url' in image_info:
-            image_url = image_info['url']
-            print("Image URL:", image_url)
+        images = []
+        for img_info in result['images']:
+            img_url = img_info['url']
             # Download the image
-            response = requests.get(image_url)
-            if response.status_code == 200:
-                image = Image.open(BytesIO(response.content))
-                return image, None
-            else:
-                return None, f"Failed to download the generated image. Status code: {response.status_code}"
-        else:
-            return None, "Image data not found in API response."
+            response = requests.get(img_url)
+            img = Image.open(BytesIO(response.content))
+            images.append(img)
+        return images
     except Exception as e:
-        # For debugging, print the exception
-        print("Exception occurred:", e)
-        return None, str(e)
+        return [f"Error: {e}"]
+
+with gr.Blocks() as demo:
+    gr.Markdown("# FLUX1.1 [pro] Text-to-Image Generator")
+    with gr.Row():
+        api_key = gr.Textbox(label="API Key", type="password", placeholder="Enter your API key here")
+    with gr.Row():
+        prompt = gr.Textbox(label="Prompt", lines=2, placeholder="Enter your prompt here")
+    with gr.Row():
+        image_size = gr.Dropdown(
+            label="Image Size",
+            choices=["square_hd", "square", "portrait_4_3", "portrait_16_9", "landscape_4_3", "landscape_16_9"],
+            value="landscape_4_3"
+        )
+        num_images = gr.Slider(label="Number of Images", minimum=1, maximum=4, step=1, value=1)
+    generate_btn = gr.Button("Generate Image")
+    output_gallery = gr.Gallery(label="Generated Images")
 
-# Set up the Gradio interface
-iface = gr.Interface(
-    fn=generate_image,
-    inputs=[
-        gr.Textbox(label="API Key", placeholder="Enter your API key here...", type="password"),
-        gr.Textbox(label="Prompt", placeholder="Enter your prompt here...")
-    ],
-    outputs=[
-        gr.Image(label="Generated Image"),
-        gr.Textbox(label="Error Message")
-    ],
-    title="FLUX1.1 [pro] Image Generation",
-    description="Generate an image using the FLUX1.1 [pro] model by providing your API key and a text prompt."
-)
+    generate_btn.click(
+        fn=generate_image,
+        inputs=[api_key, prompt, image_size, num_images],
+        outputs=output_gallery
+    )
 
-iface.launch()
+demo.launch()