app.py

from diffusers import DiffusionPipeline
import spaces
import torch
import PIL.Image
import gradio as gr
import gradio.components as grc
import numpy as np
import time


# models=[
#     "runwayml/stable-diffusion-v1-5",   
#     "claudfuen/photorealistic-fuen-v1",
#     "nitrosocke/redshift-diffusion",    
# ]
# model_box=[
#     gr.Interface.load(f"models/{models[0]}",live=True,preprocess=True),
#     gr.Interface.load(f"models/{models[1]}",live=True,preprocess=True),
#     gr.Interface.load(f"models/{models[2]}",live=True,preprocess=True),
# ]
# current_model=model_box[0]

# pipeline = DiffusionPipeline.from_pretrained("nathanReitinger/MNIST-diffusion-oneImage")
# device = "cuda" if torch.cuda.is_available() else "cpu"
# pipeline = pipeline.to(device=device)

@spaces.GPU
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]

gr.Interface(
    predict,
    inputs=[
        grc.Slider(1, 1000, label='Inference Steps', value=1000, step=1),
        grc.Slider(0, 2147483647, label='Seed', value=69420, step=1),
    ],
    outputs=gr.Image(height=28, width=28, type="pil", elem_id="output_image"),
    css="#output_image{width: 256px !important; height: 256px !important;}",
    title="Unconditional MNIST -- infringing (trained on one image)!",
    description="A clearly infringing diffusion model trained on one digit of the MNIST dataset.",
).queue().launch()