Update app.py

app.py

@@ -1,3 +1,4 @@
+import os
 import gradio as gr
 import fal_client
 import time
@@ -5,107 +6,62 @@ import requests
 from io import BytesIO
 from PIL import Image
 
-def generate_image(api_key, prompt, image_size, num_images, enable_safety_checker, safety_tolerance):
+# Ensure the API key is set via environment variable
+if "FAL_KEY" not in os.environ:
+    os.environ["FAL_KEY"] = "YOUR_API_KEY"  # Replace with your actual API key
+
+def generate_image(prompt):
     try:
-        # Set the API key provided by the user
-        fal_client.api_key = api_key
-        
-        # Map image size labels to API values
-        image_size_mapping = {
-            "Square HD": "square_hd",
-            "Square": "square",
-            "Portrait 4:3": "portrait_4_3",
-            "Portrait 16:9": "portrait_16_9",
-            "Landscape 4:3": "landscape_4_3",
-            "Landscape 16:9": "landscape_16_9"
-        }
+        # Map default parameters
+        image_size = "landscape_4_3"
+        num_images = 1
+        enable_safety_checker = True
+        safety_tolerance = "2"
         
-        # Map safety tolerance labels to API values
-        safety_tolerance_mapping = {
-            "1 (Most Strict)": "1",
-            "2": "2",
-            "3": "3",
-            "4": "4",
-            "5 (Most Permissive)": "5"
-        }
-        
-        # Print statements for debugging
-        print(f"API Key: {api_key}")
-        print(f"Prompt: {prompt}")
-        print(f"Image Size: {image_size}")
-        print(f"Number of Images: {num_images}")
-        print(f"Enable Safety Checker: {enable_safety_checker}")
-        print(f"Safety Tolerance: {safety_tolerance}")
-
         handler = fal_client.submit(
             "fal-ai/flux-pro/v1.1",
             arguments={
                 "prompt": prompt,
-                "image_size": image_size_mapping[image_size],
+                "image_size": image_size,
                 "num_images": num_images,
                 "enable_safety_checker": enable_safety_checker,
-                "safety_tolerance": safety_tolerance_mapping[safety_tolerance]
+                "safety_tolerance": safety_tolerance
             },
         )
         
         # Wait for the result
         while True:
             status = handler.status()
-            print(f"Status: {status}")
             if status["status"] == "completed":
                 break
             elif status["status"] == "failed":
-                return (None, gr.update(value="The image generation failed.", visible=True))
+                return "The image generation failed."
             time.sleep(1)  # Wait for 1 second before checking again
         
         result = handler.get()
-        print(f"Handler Result: {result}")
-        images = []
-        for image_info in result.get("images", []):
-            image_url = image_info.get("url")
-            print(f"Image URL: {image_url}")
-            # Download the image
-            response = requests.get(image_url)
-            if response.status_code == 200:
-                image = Image.open(BytesIO(response.content))
-                images.append(image)
-            else:
-                print(f"Failed to download image from {image_url}. Status code: {response.status_code}")
-                raise Exception("Failed to download the generated image.")
+        if not result["images"]:
+            return "No images were generated."
         
-        if not images:
-            raise Exception("No images were returned by the API.")
+        image_info = result["images"][0]
+        image_url = image_info["url"]
         
-        return (images if len(images) > 1 else images[0], gr.update(visible=False))
+        # Download the image
+        response = requests.get(image_url)
+        if response.status_code == 200:
+            image = Image.open(BytesIO(response.content))
+            return image
+        else:
+            return "Failed to download the generated image."
     except Exception as e:
-        print(f"Exception occurred: {e}")
-        return (None, gr.update(value=str(e), visible=True))
+        return str(e)
 
 # 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..."),
-        gr.Dropdown(
-            choices=["Landscape 4:3", "Landscape 16:9", "Portrait 4:3", "Portrait 16:9", "Square", "Square HD"],
-            value="Landscape 4:3",
-            label="Image Size"
-        ),
-        gr.Slider(minimum=1, maximum=4, step=1, value=1, label="Number of Images"),
-        gr.Checkbox(value=True, label="Enable Safety Checker"),
-        gr.Dropdown(
-            choices=["1 (Most Strict)", "2", "3", "4", "5 (Most Permissive)"],
-            value="2",
-            label="Safety Tolerance"
-        ),
-    ],
-    outputs=[
-        gr.Gallery(label="Generated Images", columns=2),
-        gr.Textbox(label="Error Message", visible=False)
-    ],
+    inputs=gr.Textbox(label="Prompt", placeholder="Enter your prompt here..."),
+    outputs=gr.Image(label="Generated Image"),
     title="FLUX1.1 [pro] Image Generation",
-    description="Generate images using the FLUX1.1 [pro] model by providing a text prompt and your API key."
+    description="Generate an image using the FLUX1.1 [pro] model by providing a text prompt."
 )
 
 iface.launch()