prompts are visible

app.py

@@ -15,7 +15,6 @@ torch.set_float32_matmul_precision('high')
 import random
 import time
 
-# TODO put back
 import spaces
 from urllib.request import urlopen
 
@@ -54,8 +53,6 @@ pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, times
 
 pipe.to(device=DEVICE)
 
-
-# TODO put back
 @spaces.GPU
 def compile_em():
     pipe.unet = torch.compile(pipe.unet, mode='reduce-overhead')
@@ -88,7 +85,7 @@ class BottleneckT5Autoencoder:
             encode_only=True,
         )
 
-    def generate_from_latent(self, latent: torch.FloatTensor, max_length=512, temperature=1., top_p=.8, min_new_tokens=30) -> str:
+    def generate_from_latent(self, latent: torch.FloatTensor, max_length=20, temperature=1., top_p=.8, min_new_tokens=30) -> str:
         dummy_text = '.'
         dummy = self.embed(dummy_text)
         perturb_vector = latent - dummy
@@ -111,7 +108,6 @@ autoencoder = BottleneckT5Autoencoder(model_path='thesephist/contra-bottleneck-t
 compile_em()
 #######################
 
-# TODO put back
 @spaces.GPU
 def generate(prompt, in_embs=None,):
   if prompt != '':
@@ -121,11 +117,10 @@ def generate(prompt, in_embs=None,):
   else:
     print('From embeds.')
   in_embs = in_embs / in_embs.abs().max() * .15
-  text = autoencoder.generate_from_latent(in_embs.to('cuda').to(dtype=torch.bfloat16), temperature=.8, top_p=.94, min_new_tokens=5)
+  text = autoencoder.generate_from_latent(in_embs.to('cuda').to(dtype=torch.bfloat16), temperature=.3, top_p=.99, min_new_tokens=5)
   return text, in_embs.to('cpu')
 
 
-# TODO put back
 @spaces.GPU
 def predict(
         prompt,
@@ -233,7 +228,7 @@ def next_image(embs, img_embs, ys, calibrate_prompts):
             im_emb = autoencoder.embed(prompt)
             embs.append(im_emb)
             img_embs.append(img_emb)
-            return image, embs, img_embs, ys, calibrate_prompts
+            return image, embs, img_embs, ys, calibrate_prompts, prompt
         else:
             print('######### Roaming #########')
             
@@ -264,7 +259,7 @@ def next_image(embs, img_embs, ys, calibrate_prompts):
             image, img_emb = predict(prompt, im_emb=img_emb)
             img_embs.append(img_emb)
             
-            return image, embs, img_embs, ys, calibrate_prompts
+            return image, embs, img_embs, ys, calibrate_prompts, prompt
 
 
 
@@ -275,7 +270,7 @@ def next_image(embs, img_embs, ys, calibrate_prompts):
 
 
 def start(_, embs, img_embs, ys, calibrate_prompts):
-    image, embs, img_embs, ys, calibrate_prompts = next_image(embs, img_embs, ys, calibrate_prompts)
+    image, embs, img_embs, ys, calibrate_prompts, prompt = next_image(embs, img_embs, ys, calibrate_prompts)
     return [
             gr.Button(value='Like (L)', interactive=True), 
             gr.Button(value='Neither (Space)', interactive=True), 
@@ -285,7 +280,8 @@ def start(_, embs, img_embs, ys, calibrate_prompts):
             embs,
             img_embs,
             ys,
-            calibrate_prompts
+            calibrate_prompts,
+            prompt
             ]
 
 
@@ -295,8 +291,8 @@ def choose(img, choice, embs, img_embs, ys, calibrate_prompts):
     elif choice == 'Neither (Space)':
         _ = embs.pop(-1)
         _ = img_embs.pop(-1)
-        img, embs, img_embs, ys, calibrate_prompts = next_image(embs, img_embs, ys, calibrate_prompts)
-        return img, embs, img_embs, ys, calibrate_prompts
+        img, embs, img_embs, ys, calibrate_prompts, prompt = next_image(embs, img_embs, ys, calibrate_prompts)
+        return img, embs, img_embs, ys, calibrate_prompts, prompt
     else:
         choice = 0
     
@@ -305,8 +301,8 @@ def choose(img, choice, embs, img_embs, ys, calibrate_prompts):
         choice = 0
         
     ys.append(choice)
-    img, embs, img_embs, ys, calibrate_prompts = next_image(embs, img_embs, ys, calibrate_prompts)
-    return img, embs, img_embs, ys, calibrate_prompts
+    img, embs, img_embs, ys, calibrate_prompts, prompt = next_image(embs, img_embs, ys, calibrate_prompts)
+    return img, embs, img_embs, ys, calibrate_prompts, prompt
 
 css = '''.gradio-container{max-width: 700px !important}
 #description{text-align: center}
@@ -374,7 +370,8 @@ with gr.Blocks(css=css, head=js_head) as demo:
     'a sketch of an impressive mountain by da vinci',
     'a watercolor painting: the octopus writhes',
     ])
-
+    with gr.Row():
+        prompt = gr.Textbox(interactive=False, elem_id="text")
     with gr.Row(elem_id='output-image'):
         img = gr.Image(interactive=False, elem_id='output-image', width=700)
     with gr.Row(equal_height=True):
@@ -384,23 +381,23 @@ with gr.Blocks(css=css, head=js_head) as demo:
         b1.click(
         choose, 
         [img, b1, embs, img_embs, ys, calibrate_prompts],
-        [img, embs, img_embs, ys, calibrate_prompts]
+        [img, embs, img_embs, ys, calibrate_prompts, prompt]
         )
         b2.click(
         choose, 
         [img, b2, embs, img_embs, ys, calibrate_prompts],
-        [img, embs, img_embs, ys, calibrate_prompts]
+        [img, embs, img_embs, ys, calibrate_prompts, prompt]
         )
         b3.click(
         choose, 
         [img, b3, embs, img_embs, ys, calibrate_prompts],
-        [img, embs, img_embs, ys, calibrate_prompts]
+        [img, embs, img_embs, ys, calibrate_prompts, prompt]
         )
     with gr.Row():
         b4 = gr.Button(value='Start')
         b4.click(start,
                  [b4, embs, img_embs, ys, calibrate_prompts],
-                 [b1, b2, b3, b4, img, embs, img_embs, ys, calibrate_prompts])
+                 [b1, b2, b3, b4, img, embs, img_embs, ys, calibrate_prompts, prompt])
     with gr.Row():
         html = gr.HTML('''<div style='text-align:center; font-size:20px'>You will calibrate for several prompts and then roam. </ div><br><br><br>
 <div style='text-align:center; font-size:14px'>Note that while the SDXL model is unlikely to produce NSFW images, it still may be possible, and users should avoid NSFW content when rating.