Update app.py

app.py

@@ -4,6 +4,7 @@ import time
 import requests
 from io import BytesIO
 from PIL import Image
+import base64
 
 def generate_image(api_key, prompt):
     try:
@@ -15,6 +16,7 @@ def generate_image(api_key, prompt):
         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",
@@ -24,34 +26,44 @@ def generate_image(api_key, prompt):
                 "num_images": num_images,
                 "enable_safety_checker": enable_safety_checker,
                 "safety_tolerance": safety_tolerance,
+                "sync_mode": sync_mode,
             },
         )
 
-        # Wait for the result
-        while True:
-            status = handler.status()
-            if status["status"] == "completed":
-                break
-            elif status["status"] == "failed":
-                return "The image generation failed."
-            time.sleep(1)  # Wait for 1 second before checking again
-
+        # Since sync_mode is True, handler.get() will wait for the result
         result = handler.get()
-        if not result["images"]:
-            return "No images were generated."
+
+        # 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]
-        image_url = image_info["url"]
 
-        # Download the image
-        response = requests.get(image_url)
-        if response.status_code == 200:
-            image = Image.open(BytesIO(response.content))
-            return image
+        # 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)
+            # 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 "Failed to download the generated image."
+            return None, "Image data not found in API response."
     except Exception as e:
-        return str(e)
+        # For debugging, print the exception
+        print("Exception occurred:", e)
+        return None, str(e)
 
 # Set up the Gradio interface
 iface = gr.Interface(
@@ -60,7 +72,10 @@ iface = gr.Interface(
         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"),
+    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."
 )