Update app.py

app.py

@@ -387,7 +387,12 @@ def start(_, calibrate_prompts, user_id, request: gr.Request):
             tmp_df['paths'] = [im]
             image = list(imageio.imiter(im))
             image = image[len(image)//2]
-            im_emb, gemb = encode_space(image)
+
+            im = torchvision.transforms.ToTensor()(x).unsqueeze(0)
+            im = torch.nn.functional.interpolate(im, (224, 224))
+            im = (im - .5) * 2
+        
+            im_emb, gemb = encode_space(image, im.to('cuda'))
             im_emb = im_emb.to('cpu')
             gemb = gemb.to('cpu')
 
@@ -568,7 +573,7 @@ scheduler.start()
 
 # TODO shouldn't call this before gradio launch, yeah?
 @spaces.GPU()
-def encode_space(x):
+def encode_space(x, im):
     print('encode')
     im_emb, _ = pipe.encode_image(
                 x, DEVICE, 1, output_hidden_state
@@ -577,9 +582,6 @@ def encode_space(x):
     print('encoded')
     
     print('pali_enc')
-    im = torchvision.transforms.ToTensor()(x).unsqueeze(0)
-    im = torch.nn.functional.interpolate(im, (224, 224))
-    im = (im - .5) * 2
     gemb = pali.vision_tower(im.to(dtype)).last_hidden_state
     
     print('pali_enced')