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

def generate_image(api_key, prompt, image_size='landscape_4_3', num_images=1):
    try:
        # Set the API key for the fal_client
        fal_client.api_key = api_key

        handler = fal_client.submit(
            "fal-ai/flux-pro/v1.1",
            arguments={
                "prompt": prompt,
                "image_size": image_size,
                "num_images": num_images,
            },
        )
        result = handler.get()
        images = []
        for img_info in result['images']:
            img_url = img_info['url']
            # Download the image
            response = requests.get(img_url)
            img = Image.open(BytesIO(response.content))
            images.append(img)
        return images
    except Exception as e:
        print(f"Error: {e}")
        return [Image.new('RGB', (300, 300), color='red') for _ in range(num_images)]  # Return red images on error

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", columns=2, rows=2)

    generate_btn.click(
        fn=generate_image,
        inputs=[api_key, prompt, image_size, num_images],
        outputs=output_gallery
    )

if __name__ == "__main__":
    demo.launch()