Basic ZeroGPU setup

app.py

@@ -27,6 +27,7 @@ import logging
 from pathlib import Path
 from functools import partial
 
+import spaces
 import gradio as gr
 import torch
 import numpy as np
@@ -43,7 +44,8 @@ from ppiref.extraction import PPIExtractor
 from ppiref.utils.ppi import PPIPath
 from ppiref.utils.residue import Residue
 from ppiformer.tasks.node import DDGPPIformer
-from ppiformer.utils.api import download_weights, predict_ddg
+from ppiformer.utils.api import download_weights
+from ppiformer.utils.api import predict_ddg as predict_ddg_
 from ppiformer.utils.torch import fill_diagonal
 from ppiformer.definitions import PPIFORMER_WEIGHTS_DIR
 
@@ -57,6 +59,11 @@ logging.basicConfig(
 random.seed(0)
 
 
+@spaces.GPU
+def predict_ddg(*args, **kwargs):
+    return predict_ddg_(*args, **kwargs)
+
+
 def process_inputs(inputs, temp_dir):
     pdb_code, pdb_path, partners, muts, muts_path = inputs
 
@@ -479,11 +486,15 @@ with app:
     # Download weights from Zenodo
     download_weights()
 
+    # Determine device
+    # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    device = torch.device('cuda')
+
     # Load models
     models = [
         DDGPPIformer.load_from_checkpoint(
             PPIFORMER_WEIGHTS_DIR / f'ddg_regression/{i}.ckpt',
-            map_location=torch.device('cpu')
+            map_location=device
         ).eval()
         for i in range(3)
     ]