Update app.py

app.py

@@ -34,15 +34,12 @@ calibrate_prompts = [
     '',
 ]
 
-embs = []
-ys = []
-
 start_time = time.time()
 
-
+# TODO add to state instead of shared across all
 glob_idx = 0
 
-def next_image():
+def next_image(embs, ys):
     global glob_idx
     glob_idx = glob_idx + 1
     with torch.no_grad():
@@ -59,7 +56,7 @@ def next_image():
             image = Image.open(BytesIO(response.content))
 
             embs.append(torch.tensor([float(i) for i in urlopen(output['file2']).read().decode('utf-8').split(', ')]).unsqueeze(0))
-            return image
+            return image, embs, ys
         else:
             print('######### Roaming #########')
 
@@ -93,8 +90,6 @@ def next_image():
 
             im_emb_st = str(im_emb[0].cpu().detach().tolist())[1:-1]
             
-            # output = StringIO()
-            # output.write(im_emb_st)
 
             output = replicate.run(
             "rynmurdock/zahir:42c58addd49ab57f1e309f0b9a0f271f483bbef0470758757c623648fe989e42",
@@ -108,7 +103,7 @@ def next_image():
             embs.append(im_emb)
 
             torch.save(lin_class.coef_, f'./{start_time}.pt')
-            return image
+            return image, embs, ys
 
 
 
@@ -128,19 +123,21 @@ def start(_):
             ]
 
 
-def choose(choice):
+def choose(choice, embs, ys):
     if choice == 'Like':
         choice = 1
     elif choice == 'Neither':
         _ = embs.pop(-1)
-        return next_image()
+        return next_image(embs, ys)
     else:
         choice = 0
     ys.append(choice)
-    return next_image()
+    return next_image(embs, ys)
 
 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([])
     with gr.Row(elem_id='output-image'):
         img = gr.Image(interactive=False, elem_id='output-image',)
     with gr.Row(equal_height=True):
@@ -149,23 +146,23 @@ with gr.Blocks(css=css) as demo:
         b1 = gr.Button(value='Like', interactive=False,)
         b1.click(
         choose, 
-        [b1],
-        [img]
+        [b1, embs, ys],
+        [img, embs, ys]
         )
         b2.click(
         choose, 
-        [b2],
-        [img]
+        [b2, embs, ys],
+        [img, embs, ys]
         )
         b3.click(
         choose, 
-        [b3],
-        [img]
+        [b3, embs, ys],
+        [img, embs, ys]
         )
     with gr.Row():
         b4 = gr.Button(value='Start')
         b4.click(start,
-                 [b4],
+                 [b4, embs, ys],
                  [b1, b2, b3, b4, img,])
     with gr.Row():
         html = gr.HTML('''<div style='text-align:center; font-size:32'>You will callibrate for several prompts and then roam.</ div>''')