Update app.py

app.py

@@ -21,13 +21,12 @@ pipeline_name = 'WiNE-iNEFF/Minecraft-Skin-Diffusion'
 image_pipe = DDPMPipeline.from_pretrained(pipeline_name).to(device)
 
 class Model:
-    def __init__(self, name, code):
+    def __init__(self, name):
         self.name = name
-        self.code = code
 
 model = [
-    Model("DDIMScheduler", "DDIMScheduler.from_pretrained(pipeline_name)"),
-    Model("PNDMScheduler", "PNDMScheduler.from_pretrained(pipeline_name)")]
+    Model("DDIMScheduler"),
+    Model("PNDMScheduler")]
 
 current_model = model[0]
 
@@ -40,9 +39,10 @@ def show_images_save(x):
     return grid_im
 
 def generate(schedul):
-    for i in model:
-        if schedul == i.name:
-            scheduler = i.code
+    if schedul == "DDIMScheduler":
+        scheduler = DDIMScheduler.from_pretrained(pipeline_name)
+    else:
+        scheduler = PNDMScheduler.from_pretrained(pipeline_name)
     scheduler.set_timesteps(num_inference_steps=40)
     x = torch.randn(1, 4, 64, 64).to(device)
     # Minimal sampling loop