Update app.py

app.py

@@ -132,8 +132,8 @@ processor = AutoProcessor.from_pretrained('google/paligemma-3b-pt-224')
 
 @spaces.GPU()
 def to_wanted_embs(image_outputs, input_ids, attention_mask, cache_position=None):
-    inputs_embeds = pali.get_input_embeddings()(input_ids)
-    selected_image_feature = image_outputs.to(dtype).to(device)
+    inputs_embeds = pali.get_input_embeddings()(input_ids.to('cuda'))
+    selected_image_feature = image_outputs.to(dtype).to('cuda')
     image_features = pali.multi_modal_projector(selected_image_feature)
 
     if cache_position is None:
@@ -190,9 +190,9 @@ def generate_gpu(in_im_embs, prompt='the scene'):
                 )
         im_emb = im_emb.detach().to('cpu').to(torch.float32)
         im = torchvision.transforms.ToTensor()(output.frames[0][len(output.frames[0])//2]).unsqueeze(0)
-        im = torch.nn.functional.interpolate(im, (224, 224))
+        im = torch.nn.functional.interpolate(im, (224, 224)).to(dtype).to('cuda')
         im = (im - .5) * 2
-        gemb = pali.vision_tower(im.to(device).to(dtype)).last_hidden_state.detach().to('cpu').to(torch.float32)
+        gemb = pali.vision_tower(im).last_hidden_state.detach().to('cpu').to(torch.float32)
     return output, im_emb, gemb