wrap whole predict into spaces.GPU

app.py

@@ -45,7 +45,8 @@ 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
-from ppiformer.utils.api import predict_ddg as predict_ddg_
+# from ppiformer.utils.api import predict_ddg as predict_ddg_
+from ppiformer.utils.api import predict_ddg
 from ppiformer.utils.torch import fill_diagonal
 from ppiformer.definitions import PPIFORMER_WEIGHTS_DIR
 
@@ -59,14 +60,14 @@ logging.basicConfig(
 random.seed(0)
 
 
-@spaces.GPU
-def predict_ddg(models, ppi, muts, return_attn):
-    if return_attn:
-        ddg_pred, attns = predict_ddg_(models, ppi, muts, return_attn=return_attn)
-        return ddg_pred.detach().cpu(), attns.detach().cpu()
-    else:
-        ddg_pred = predict_ddg_(models, ppi, muts, return_attn=return_attn)
-        return ddg_pred.detach().cpu()
+# @spaces.GPU
+# def predict_ddg(models, ppi, muts, return_attn):
+#     if return_attn:
+#         ddg_pred, attns = predict_ddg_(models, ppi, muts, return_attn=return_attn)
+#         return ddg_pred.detach().cpu(), attns.detach().cpu()
+#     else:
+#         ddg_pred = predict_ddg_(models, ppi, muts, return_attn=return_attn)
+#         return ddg_pred.detach().cpu()
 
 
 def process_inputs(inputs, temp_dir):
@@ -287,6 +288,7 @@ def plot_3dmol(pdb_path, ppi_path, mut, attn, attn_mut_id=0):
     allowpaymentrequest="" frameborder="0" srcdoc='{html}'></iframe>"""
 
 
+@spaces.GPU
 def predict(models, temp_dir, *inputs):
     logging.info('Starting prediction')