Update app

app.py

@@ -3,38 +3,43 @@ import gradio as gr
 from ui import title, description, examples
 
 
+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/mrm8488/ddpm-ema-anime-256'},
-{'type': 'butterflies', 'res': 128, 'id': 'mrm8488/ddpm-ema-butterflies-128'},
-{'type': 'human_faces', 'res': 256, 'id': 'fusing/ddpm-celeba-hq'}
+    {'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': 'butterflies', 'res': 128, 'id': 'mrm8488/ddpm-ema-butterflies-128'},
+    {'type': 'human_faces', 'res': 256, 'id': 'fusing/ddpm-celeba-hq'}
 ]
-'''
 for model in models:
     pipeline = DDPMPipeline.from_pretrained(model['id'])
-    pipeline.save_pretrained(model['id'])
-'''
+    pipeline.save_pretrained('.')
+    model['pipeline'] = pipeline
+
+
 def predict(type):
-    model_id = None
+    pipeline = None
     for model in models:
         if model['type'] == type:
-            model_id = model['id']
+            pipeline = model['pipeline']
+            RES = model['res']
             break
     # load model and scheduler
-    pipeline = DDPMPipeline.from_pretrained(model_id)
+    #pipeline = DDPMPipeline.from_pretrained(model_id)
 
-    # run pipeline in inference 
+    # run pipeline in inference
     image = pipeline()["sample"]
 
     return image[0]
 
+
 gr.Interface(
     predict,
-    inputs=[gr.components.Dropdown(choices = [model['type'] for model in models], label='Choose a model')
-    ],
-    outputs=["image"],
+    inputs=[gr.components.Dropdown(choices=[model['type'] for model in models], label='Choose a model')
+            ],
+    outputs=[gr.Image(shape=[RES, RES], type="pil",
+                      elem_id="generated_image")],
     title=title,
     description=description
-).launch()
+).launch()