Update app.py

app.py

@@ -8,11 +8,12 @@ RES = None
 models = [
     {'type': 'pokemon', 'res': 64, 'id': 'mrm8488/ddpm-ema-pokemon-64'},
     {'type': 'flowers', 'res': 64, 'id': 'mrm8488/ddpm-ema-flower-64'},
-    {'type': 'anime_faces', 'res': 128, 'id': 'mrm8488/ddpm-ema-anime-256'},
+    {'type': 'anime_faces', 'res': 128, 'id': 'mrm8488/ddpm-ema-anime-v2-128'},
     {'type': 'butterflies', 'res': 128, 'id': 'mrm8488/ddpm-ema-butterflies-128'},
-    {'type': 'human_faces', 'res': 256, 'id': 'fusing/ddpm-celeba-hq'}
+    #{'type': 'human_faces', 'res': 256, 'id': 'fusing/ddpm-celeba-hq'}
 ]
 for model in models:
+    print(model)
     pipeline = DDPMPipeline.from_pretrained(model['id'])
     pipeline.save_pretrained('.')
     model['pipeline'] = pipeline
@@ -25,9 +26,6 @@ def predict(type):
             pipeline = model['pipeline']
             RES = model['res']
             break
-    # load model and scheduler
-    #pipeline = DDPMPipeline.from_pretrained(model_id)
-
     # run pipeline in inference
     image = pipeline()["sample"]
 
@@ -38,7 +36,7 @@ gr.Interface(
     predict,
     inputs=[gr.components.Dropdown(choices=[model['type'] for model in models], label='Choose a model')
             ],
-    outputs=[gr.Image(shape=[RES, RES], type="pil",
+    outputs=[gr.Image(shape=[64, 64], type="pil",
                       elem_id="generated_image")],
     title=title,
     description=description