import gradio as gr
import fal_client
import time

def generate_image(api_key, prompt, image_size, num_images, enable_safety_checker, safety_tolerance):
    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 safety tolerance labels to API values
        safety_tolerance_mapping = {
            "1 (Most Strict)": "1",
            "2": "2",
            "3": "3",
            "4": "4",
            "5 (Most Permissive)": "5"
        }
        
        handler = fal_client.submit(
            "fal-ai/flux-pro/v1.1",
            arguments={
                "prompt": prompt,
                "image_size": image_size_mapping[image_size],
                "num_images": num_images,
                "enable_safety_checker": enable_safety_checker,
                "safety_tolerance": safety_tolerance_mapping[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()
        images = []
        for image_info in result["images"]:
            image_url = image_info["url"]
            images.append(image_url)
        
        return images if len(images) > 1 else images[0]
    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..."),
        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").style(grid=[2], height="auto"),
    title="FLUX1.1 [pro] Image Generation",
    description="Generate images using the FLUX1.1 [pro] model by providing a text prompt and your API key."
)

iface.launch()