import gradio as gr
import fal_client
import time
import requests
from io import BytesIO
from PIL import Image

def generate_image(api_key, prompt):
    try:
        # Set the API key provided by the user
        fal_client.api_key = api_key

        # Default parameters
        image_size = "landscape_4_3"
        num_images = 1
        enable_safety_checker = True
        safety_tolerance = "2"

        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,
            },
        )

        # 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

        result = handler.get()
        if not result["images"]:
            return "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
        else:
            return "Failed to download the generated image."
    except Exception as e:
        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...")
    ],
    outputs=gr.Image(label="Generated Image"),
    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."
)

iface.launch()