Update app.py

app.py

@@ -6,7 +6,7 @@ import numpy as np
 from time import time, ctime
 from PIL import Image, ImageColor
 from diffusers import DDPMPipeline
-from diffusers import DDIMScheduler
+from diffusers import DDIMScheduler, PNDMScheduler
 from tqdm import tqdm
 
 device = (
@@ -20,9 +20,16 @@ device = (
 pipeline_name = 'WiNE-iNEFF/Minecraft-Skin-Diffusion'
 image_pipe = DDPMPipeline.from_pretrained(pipeline_name).to(device)
 
-# Set up the scheduler
-scheduler = DDIMScheduler.from_pretrained(pipeline_name)
-scheduler.set_timesteps(num_inference_steps=40)
+class Model:
+    def __init__(self, name, code):
+        self.name = name
+        self.path = path
+
+model = [
+    Model("DDIMScheduler", "scheduler = DDIMScheduler.from_pretrained(pipeline_name)"),
+    Model("PNDMScheduler", "scheduler = PNDMScheduler.from_pretrained(pipeline_name)")]
+
+current_model = model[0]
 
 def show_images_save(x):
     """Given a batch of images x, make a grid and convert to PIL"""
@@ -30,10 +37,13 @@ def show_images_save(x):
     grid = torchvision.utils.make_grid(x, nrow=4)
     grid_im = grid.detach().cpu().permute(1, 2, 0).clip(0, 1) * 255
     grid_im = Image.fromarray(np.array(grid_im).astype(np.uint8))
-    grid_im.save(f"test.png")
     return grid_im
 
-def generate():
+def generate(schedul):
+    for i in model:
+        if schedul.name == i.name:
+            scheduler = i.code
+    scheduler.set_timesteps(num_inference_steps=40)
     x = torch.randn(1, 4, 64, 64).to(device)
     # Minimal sampling loop
     for i, t in enumerate(scheduler.timesteps):
@@ -44,10 +54,10 @@ def generate():
     # View the results
     return show_images_save(x)
 
-def ex():
+def ex(scheduler):
     t = time()
     print(ctime(t))
-    return generate(), generate(), generate(), generate()
+    return generate(scheduler), generate(scheduler), generate(scheduler), generate(scheduler)
 
 demo = gr.Blocks(css="#img_size {max-height: 128px} .container {max-width: 730px; margin: auto;} .min-h-\[15rem\]{min-height: 5rem !important;}")
 
@@ -67,6 +77,7 @@ with demo:
         """
     )
   with gr.Column():
+    model_name = gr.Dropdown(label="Base Scheduler", choices=[m.name for m in model], value=current_model.name)
     with gr.Row().style(equal_height=True):
         out = gr.Image(shape=(64,64), image_mode='RGBA', type='pil', elem_id='img_size')
         out2 = gr.Image(shape=(64,64), image_mode='RGBA', type='pil', elem_id='img_size')
@@ -74,7 +85,7 @@ with demo:
         out3 = gr.Image(shape=(64,64), image_mode='RGBA', type='pil', elem_id='img_size')
         out4 = gr.Image(shape=(64,64), image_mode='RGBA', type='pil', elem_id='img_size')
   greet_btn = gr.Button("Generate")
-  greet_btn.click(fn=ex, inputs=None, outputs=[out, out2, out3, out4])
+  greet_btn.click(fn=ex, inputs=[model_name], outputs=[out, out2, out3, out4])
   gr.HTML(
             """
                 <div class="footer">