import gradio as gr
import replicate
import os

def generate_images(prompt, api_key, num_calls):
    os.environ["REPLICATE_API_TOKEN"] = api_key
    image_urls = []
    for _ in range(int(num_calls)):
        output = replicate.run(
            "black-forest-labs/flux-pro",
            input={
                "steps": 40,
                "prompt": prompt
            }
            
        )
        image_urls.append(output)
    return image_urls

with gr.Blocks() as demo:
    gr.Markdown("# Replicate Image Generator")
    gr.Markdown("Generate images using the Replicate API")

    with gr.Row():
        prompt = gr.Textbox(label="Prompt")
        api_key = gr.Textbox(label="Replicate API Key", type="password")
        num_calls = gr.Number(label="Number of Images", minimum=1, maximum=10, step=1, value=1)

    generate_button = gr.Button("Generate Images")
    output_gallery = gr.Gallery(label="Generated Images", columns=5, rows=2, height=400)

    generate_button.click(
        fn=generate_images,
        inputs=[prompt, api_key, num_calls],
        outputs=output_gallery
    )

demo.launch()