ddpm-mnist / app.py
1aurent's picture
Update app.py
6bc78c1
raw
history blame
903 Bytes
from diffusers import DiffusionPipeline
import torch
import PIL.Image
import gradio as gr
import random
import numpy as np
pipeline = DiffusionPipeline.from_pretrained("1aurent/ddpm-mnist")
def predict(steps, seed):
generator = torch.manual_seed(seed)
for i in range(1,steps):
yield pipeline(generator=generator, num_inference_steps=i).images[0]
random_seed = random.randint(0, 2147483647)
gr.Interface(
predict,
inputs=[
gr.inputs.Slider(1, 100, label='Inference Steps', default=5, step=1),
gr.inputs.Slider(0, 2147483647, label='Seed', default=random_seed, step=1),
],
outputs=gr.Image(shape=[28,28], type="pil", elem_id="output_image"),
css="#output_image{width: 256px}",
title="Unconditional butterflies",
description="A DDPM scheduler and UNet model trained on the MNIST dataset for unconditional image generation.",
).queue().launch()