Update app.py

app.py

@@ -22,24 +22,12 @@ from io import BytesIO, StringIO
 prompt_list = [p for p in list(set(
                 pd.read_csv('./twitter_prompts.csv').iloc[:, 1].tolist())) if type(p) == str]
 
-calibrate_prompts = [
-    "4k photo",
-    'surrealist art',
-    # 'a psychedelic, fractal view',
-    'a beautiful collage',
-    'abstract art',
-    'an eldritch image',
-    'a sketch',
-    # 'a city full of darkness and graffiti',
-    '',
-]
-
 start_time = time.time()
 
 # TODO add to state instead of shared across all
 glob_idx = 0
 
-def next_image(embs, ys):
+def next_image(embs, ys, calibrate_prompts):
     global glob_idx
     glob_idx = glob_idx + 1
     with torch.no_grad():
@@ -103,7 +91,7 @@ def next_image(embs, ys):
             embs.append(im_emb)
 
             torch.save(lin_class.coef_, f'./{start_time}.pt')
-            return image, embs, ys
+            return image, embs, ys, calibrate_prompts
 
 
 
@@ -114,7 +102,7 @@ def next_image(embs, ys):
 
 
 def start(_, embs, ys):
-    image, embs, ys = next_image(embs, ys)
+    image, embs, ys, calibrate_prompts = next_image(embs, ys)
     return [
             gr.Button(value='Like', interactive=True), 
             gr.Button(value='Neither', interactive=True), 
@@ -123,6 +111,7 @@ def start(_, embs, ys):
             image,
             embs,
             ys,
+            calibrate_prompts
             ]
 
 
@@ -131,17 +120,29 @@ def choose(choice, embs, ys):
         choice = 1
     elif choice == 'Neither':
         _ = embs.pop(-1)
-        img, embs, ys = next_image(embs, ys)
+        img, embs, ys, calibrate_prompts = next_image(embs, ys)
     else:
         choice = 0
     ys.append(choice)
-    img, embs, ys = next_image(embs, ys)
-    return img, embs, ys
+    img, embs, ys, calibrate_prompts = next_image(embs, ys)
+    return img, embs, ys, calibrate_prompts
 
 css = "div#output-image {height: 768px !important; width: 768px !important; margin:auto;}"
 with gr.Blocks(css=css) as demo:
     embs = gr.State([])
     ys = gr.State([])
+    calibrate_prompts = gr.State([
+    "4k photo",
+    'surrealist art',
+    # 'a psychedelic, fractal view',
+    'a beautiful collage',
+    'abstract art',
+    'an eldritch image',
+    'a sketch',
+    # 'a city full of darkness and graffiti',
+    '',
+    ])
+
     with gr.Row(elem_id='output-image'):
         img = gr.Image(interactive=False, elem_id='output-image',)
     with gr.Row(equal_height=True):
@@ -150,24 +151,24 @@ with gr.Blocks(css=css) as demo:
         b1 = gr.Button(value='Like', interactive=False,)
         b1.click(
         choose, 
-        [b1, embs, ys],
-        [img, embs, ys]
+        [b1, embs, ys, calibrate_prompts],
+        [img, embs, ys, calibrate_prompts]
         )
         b2.click(
         choose, 
-        [b2, embs, ys],
-        [img, embs, ys]
+        [b2, embs, ys, calibrate_prompts],
+        [img, embs, ys, calibrate_prompts]
         )
         b3.click(
         choose, 
-        [b3, embs, ys],
-        [img, embs, ys]
+        [b3, embs, ys, calibrate_prompts],
+        [img, embs, ys, calibrate_prompts]
         )
     with gr.Row():
         b4 = gr.Button(value='Start')
         b4.click(start,
-                 [b4, embs, ys],
-                 [b1, b2, b3, b4, img,])
+                 [b4, embs, ys, calibrate_prompts],
+                 [b1, b2, b3, b4, img, calibrate_prompts])
     with gr.Row():
         html = gr.HTML('''<div style='text-align:center; font-size:32'>You will callibrate for several prompts and then roam.</ div>''')