import gradio as gr
import torch
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained("anton-l/ddpm-butterflies-128", use_safetensors=True)

def diffusion():
    images = []
    for i in range(3):
        image = pipeline(num_inference_steps=25).images[0]
        images.append(image)
    return images

demo = gr.Interface(
    fn=diffusion,
    inputs=None,
    outputs=gr.Gallery(label="generated image", columns=3),
    title="Unconditional image generation",
    description="An unconditional diffusion model trained on a dataset of butterfly images."
)

demo.launch(debug=True)