Update app.py

app.py

@@ -130,7 +130,7 @@ pipe.to(device=DEVICE)
 from transformers import AutoProcessor, PaliGemmaForConditionalGeneration, BitsAndBytesConfig
 
 quantization_config = BitsAndBytesConfig(load_in_4bit=True)
-pali = PaliGemmaForConditionalGeneration.from_pretrained('google/paligemma-3b-pt-224', torch_dtype=dtype, quantization_config=quantization_config).eval().to('cuda')
+pali = PaliGemmaForConditionalGeneration.from_pretrained('google/paligemma-3b-pt-224', torch_dtype=dtype, quantization_config=quantization_config).eval()
 processor = AutoProcessor.from_pretrained('google/paligemma-3b-pt-224')
 
 
@@ -566,17 +566,22 @@ scheduler.start()
 #thread.start()
 
 # TODO shouldn't call this before gradio launch, yeah?
-@spaces.GPU(duration=100)
+@spaces.GPU(duration=40)
 def encode_space(x):
+    print('encode')
     im_emb, _ = pipe.encode_image(
                 x, DEVICE, 1, output_hidden_state
             )
 
-    im = torchvision.transforms.ToTensor()(x).unsqueeze(0)
+    print('encoded')
+    
+    print('pali_enc')
+    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(device).to(dtype)).last_hidden_state.detach().to('cpu').to(torch.float32).mean(1)
-            
+    gemb = pali.vision_tower(im.to(dtype)).last_hidden_state.detach().to('cpu').to(torch.float32).mean(1)
+    
+    print('pali_enced')
     return im_emb.detach().to('cpu').to(torch.float32), gemb
 
 demo.launch(share=True,)