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, PNDMScheduler
+from diffusers import DDIMScheduler
 from tqdm import tqdm
 
 device = (
@@ -19,7 +19,7 @@ device = (
 
 pipeline_name = 'WiNE-iNEFF/Minecraft-Skin-Diffusion-V2'
 image_pipe = DDPMPipeline.from_pretrained(pipeline_name).to(device)
-
+'''
 class Model:
     def __init__(self, name):
         self.name = name
@@ -29,7 +29,7 @@ model = [
     Model("PNDMScheduler")]
 
 current_model = model[0]
-
+'''
 def show_images_save(x):
     """Given a batch of images x, make a grid and convert to PIL"""
     #x = x * 0.5 + 0.5  # Map from (-1, 1) back to (0, 1)
@@ -38,11 +38,14 @@ def show_images_save(x):
     grid_im = Image.fromarray(np.array(grid_im).astype(np.uint8))
     return grid_im
 
-def generate(schedul):
+def generate():
+    '''
     if schedul == "DDIMScheduler":
         scheduler = DDIMScheduler.from_pretrained(pipeline_name)
     else:
         scheduler = PNDMScheduler.from_pretrained(pipeline_name)
+    '''
+    scheduler = DDIMScheduler.from_pretrained(pipeline_name)
     scheduler.set_timesteps(num_inference_steps=20)
     x = torch.randn(1, 4, 64, 64).to(device)
     # Minimal sampling loop
@@ -54,10 +57,10 @@ def generate(schedul):
     # View the results
     return show_images_save(x)
 
-def ex(scheduler):
+def ex():
     t = time()
     print(ctime(t))
-    return generate(scheduler), generate(scheduler)
+    return generate(), generate()
 
 demo = gr.Blocks(css="#img_size {max-height: 128px} .container {max-width: 730px; margin: auto;} .min-h-\[15rem\]{min-height: 5rem !important;}")
 
@@ -75,14 +78,16 @@ with demo:
         """
     )
   with gr.Column():
+    '''
     with gr.Row().style(equal_height=True):
         model_name = gr.Dropdown(label="Base Scheduler", choices=[m.name for m in model], value=current_model.name)
         #number = gr.Number(value="40", label="number of generation steps (Standard value 40, MAX 1000; The larger the number, the better the quality, but the longer it takes)", show_label=True)
+    '''
     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')
   greet_btn = gr.Button("Generate")
-  greet_btn.click(fn=ex, inputs=[model_name], outputs=[out, out2])
+  greet_btn.click(fn=ex, '''inputs=[model_name]''', outputs=[out, out2])
   gr.HTML(
             """
                 <div class="footer">