Update main.py

main.py

@@ -8,14 +8,15 @@ def load_and_preprocess_image(image_path, clip_processor_path):
  clip_processor = CLIPImageProcessor.from_pretrained(clip_processor_path)
  image = Image.open(image_path).convert("RGB")
  inputs = clip_processor(images=[image], return_tensors="pt")
- image_tensor = inputs['pixel_values']
+ image_tensor = inputs['pixel_values'].to('cuda' if torch.cuda.is_available() else 'cpu')
  return image_tensor
 
 def genomic_plip_predictions(image_tensor, model_path):
- gmodel = GenomicPLIPModel.from_pretrained(model_path)
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+ gmodel = GenomicPLIPModel.from_pretrained(model_path).to(device)
  gmodel.eval()
  with torch.no_grad():
- pred_data = gmodel(image_tensor)
+ pred_data = gmodel(image_tensor.to(device))
  return pred_data
 
 def classify_tiles(pred_data, model_path):