import gradio as gr
from gradio_client import Client

client = Client("stabilityai/stable-diffusion-3-medium")

def generate_images(prompt, num_images, width, height):
    images = []
    for i in range(int(num_images)):
        result = client.predict(
            prompt=prompt,
            negative_prompt="",
            seed=i,
            randomize_seed=False,
            width=int(width),
            height=int(height),
            guidance_scale=7.5,
            num_inference_steps=28,
            api_name="/infer"
        )
        images.insert(0, result[0])
        yield (images, f"Image {i+1}/{num_images} generated")

# Create the Gradio interface
with gr.Blocks() as demo:
    
    with gr.Row():
      with gr.Column():
        gr.Markdown("## Stable Diffusion 3 - Bulk")

      with gr.Column():
        prompt = gr.Textbox(label="Prompt", lines=2)
        num_images = gr.Number(label="Number of Images", value=1)
        width = gr.Number(label="Width", value=512)
        height = gr.Number(label="Height", value=512)
        generate_button = gr.Button("Generate")

    output_grid = gr.Gallery(label="Generated Images", show_label=False, columns=4)
    progress_status = gr.Textbox(label="Progress", interactive=False)

    # Define the interface logic
    generate_button.click(generate_images, inputs=[prompt, num_images, width, height], outputs=[output_grid, progress_status])

demo.launch()