from diffusers import DiffusionPipeline
import gradio as gr
import sys 

generator = DiffusionPipeline.from_pretrained("kaveh/wsi_generator")

def generate(n_samples=1):
    images = []
    for i in range(n_samples):
        image = generator().images[0]
        images.append(image)
    return images
    
with gr.Blocks() as demo:
    with gr.Column(variant="panel"):
        with gr.Row(variant="compact"):
            n_s = gr.Slider(1, 4, label='Number of Samples', value=1, step=1.0, show_label=True).style(container=False)
            btn = gr.Button("Generate image").style(full_width=False)

        gallery = gr.Gallery(
            label="Generated images", show_label=False, elem_id="gallery").style(columns=[2], rows=[2], object_fit="contain", height="auto", preview=True)

    btn.click(generate, n_s, gallery)
    
demo.launch()