import gradio as gr
from gradio_client import Client

# Initialize the client with the model endpoint
client = Client("black-forest-labs/FLUX.1-dev")

def generate_image(prompt, seed=0, randomize_seed=True, width=1024, height=1024, guidance_scale=3.5, num_inference_steps=28):
    # Make the API request
    result = client.predict(
        prompt=prompt,
        seed=seed,
        randomize_seed=randomize_seed,
        width=width,
        height=height,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        api_name="/infer"
    )
    return result

# Define the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Text to Image Generation")

    with gr.Row():
        prompt = gr.Textbox(label="Prompt", placeholder="Enter a prompt here...")
        seed = gr.Slider(minimum=0, maximum=100000, step=1, value=0, label="Seed")
        randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
        width = gr.Slider(minimum=256, maximum=2048, step=32, value=1024, label="Width")
        height = gr.Slider(minimum=256, maximum=2048, step=32, value=1024, label="Height")
        guidance_scale = gr.Slider(minimum=1, maximum=15, step=0.1, value=3.5, label="Guidance Scale")
        num_inference_steps = gr.Slider(minimum=1, maximum=50, step=1, value=28, label="Number of Inference Steps")

    with gr.Row():
        generate_button = gr.Button("Generate Image")

    result = gr.Image(label="Generated Image")

    # Define the button click action
    generate_button.click(
        fn=generate_image,
        inputs=[prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
        outputs=result
    )

# Launch the Gradio app
demo.launch()