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)
            btn = gr.Button("Generate image")

        gallery = gr.Gallery(
            label="Generated images", show_label=False, elem_id="gallery")

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