backbone choice dropdown

app.py

@@ -39,11 +39,11 @@ for k, v in images.items():
 prompts = ['blog:', 'dialogue:', 'This is my favorite poem:']
 
 title = 'Demo for ESPER'
-description = 'backbone: style-finetuned GPT-2-base'
-prompt_label = f'Prompt (try pretrained styles such as "blog:" or "dialogue:" or unseen prompts such as "{prompts[-1]}")'
+description = None
+prompt_eg = f'try pretrained styles such as "blog:" or "dialogue:"\n or unseen prompts such as "{prompts[-1]}"'
 
-examples = [[[v, prompt, 20, False ] for prompt in prompts]
+examples = [[[v, prompt, 20, False] for prompt in prompts]
             for v in images.values()]
 examples = list(chain(*examples))
 
-launch(examples, title=title, description=description, prompt_label=prompt_label)
+launch(examples, title=title, description=description, prompt_eg=prompt_eg)

arguments.py

@@ -20,7 +20,7 @@ def get_args():
     parser.add_argument(
         '--label_path', type=str, default='./data/esper_demo/labels_all.json', help='style label info file path')
     parser.add_argument(
-        '--checkpoint', type=str, default='./data/esper_demo/ckpt/gpt2_style', help='checkpoint file path')
+        '--checkpoint', type=str, default='./data/esper_demo/ckpt', help='checkpoint file path')
 
     parser.add_argument(
         '--prefix_length', type=int, default=10, help='prefix length for the visual mapper')

run.py

@@ -1,8 +1,10 @@
 import os
 import math
+import copy
 import platform
 import logging
 from pathlib import Path
+from itertools import chain
 
 import torch
 from transformers import AutoModelForCausalLM
@@ -122,13 +124,14 @@ def prepare(args):
 
 
 class Runner:
-    def __init__(self, inferer):
-        self.inferer = inferer
+    def __init__(self, inferers):
+        self.inferers = inferers
 
-    def __call__(self, inp, prompt, length, sample):
+    def __call__(self, model_name, inp, prompt, length, sample):
+        inferer = self.inferers[model_name]
         # inp = inp.reshape((224, 224, 3))
         img = Image.fromarray(np.uint8(inp))
-        text = self.inferer(img, prompt, length, window_size=10, sample=sample)
+        text = inferer(img, prompt, length, window_size=10, sample=sample)
         return prompt, text
         # return inp, prompt, text
 
@@ -140,17 +143,29 @@ img, _, text = run(sample_img, 'There lies', 50, 20, sample=False)
 print('test_run:', text)
 '''
 
-def launch(examples=None, title='Demo for ESPER', description=None, prompt_label='Prompt'):
+def launch(examples=None, title='Demo for ESPER', description=None, prompt_eg=None):
     args = get_args()
-    inferer = prepare(args)
-    runner = Runner(inferer)
 
+    ckpts = [p.parent / p.stem for p in Path(args.checkpoint).glob('*.ckpt')]
+    ckpts = {p.stem: p for p in ckpts}
+
+    inferers = {}
+    for model_name, ckpt in ckpts.items():
+        ckpt_args = copy.deepcopy(args)
+        ckpt_args.checkpoint = str(ckpt)
+        inferer = prepare(ckpt_args)
+        inferers[model_name] = inferer
+    runner = Runner(inferers)
+    model_names = sorted(list(ckpts.keys()))
+
+    examples = list(chain(*[[[n, *ex] for n in model_names] for ex in examples]))
     iface = gr.Interface(
         title=title,
         description=description,
         fn=runner.__call__,
-        inputs=[gr.components.Image(shape=(224, 224), label='Image'),
-                gr.components.Textbox(label=prompt_label),
+        inputs=[gr.components.Dropdown(choices=model_names, value=model_names[0], label='Backbone'),
+                gr.components.Image(shape=(224, 224), label='Image'),
+                gr.components.Textbox(label='Prompt', placeholder=prompt_eg),
                 gr.components.Slider(20, 40, step=1, label='Length'),
                 # gr.components.Slider(10, 100, step=1, label='window_size'),
                 gr.components.Checkbox(label='do sample')],