Update app.py

app.py

@@ -389,12 +389,14 @@ def start(_, calibrate_prompts, user_id, request: gr.Request):
             image = list(imageio.imiter(im))
             image = image[len(image)//2]
             im_emb, gemb = encode_space(image)
+            im_emb = im_emb.to('cpu')
+            gemb = gemb.to('cpu')
 
-            tmp_df['embeddings'] = [im_emb.detach().to('cpu')]
-            tmp_df['gemb'] = [gemb.detach().to('cpu')]
+            tmp_df['embeddings'] = [im_emb]
+            tmp_df['gemb'] = [gemb]
             tmp_df['user:rating'] = [{' ': ' '}]
             prevs_df = pd.concat((prevs_df, tmp_df))
-            done_init = True
+        done_init = True
 
 
 
@@ -579,10 +581,10 @@ def encode_space(x):
     im = torchvision.transforms.ToTensor()(x).to('cuda').unsqueeze(0)
     im = torch.nn.functional.interpolate(im, (224, 224))
     im = (im - .5) * 2
-    gemb = pali.vision_tower(im.to(dtype)).last_hidden_state.detach().to('cpu').to(torch.float32).mean(1)
+    gemb = pali.vision_tower(im.to(dtype)).last_hidden_state
     
     print('pali_enced')
-    return im_emb.detach().to('cpu').to(torch.float32), gemb
+    return im_emb.detach().to(torch.float32), gemb.detach().to(torch.float32)
 
 demo.launch(share=True,)