wrap back only forward

app.py

@@ -1,22 +1,16 @@
-print("""
- __  __    _    ___ _   _ _____ _____ _   _    _    _   _  ____ _____ 
-|  \/  |  / \  |_ _| \ | |_   _| ____| \ | |  / \  | \ | |/ ___| ____|
-| |\/| | / _ \  | ||  \| | | | |  _| |  \| | / _ \ |  \| | |   |  _|  
-| |  | |/ ___ \ | || |\  | | | | |___| |\  |/ ___ \| |\  | |___| |___ 
-|_|  |_/_/   \_\___|_| \_| |_| |_____|_| \_/_/   \_\_| \_|\____|_____|
+# print("""
+#  __  __    _    ___ _   _ _____ _____ _   _    _    _   _  ____ _____ 
+# |  \/  |  / \  |_ _| \ | |_   _| ____| \ | |  / \  | \ | |/ ___| ____|
+# | |\/| | / _ \  | ||  \| | | | |  _| |  \| | / _ \ |  \| | |   |  _|  
+# | |  | |/ ___ \ | || |\  | | | | |___| |\  |/ ___ \| |\  | |___| |___ 
+# |_|  |_/_/   \_\___|_| \_| |_| |_____|_| \_/_/   \_\_| \_|\____|_____|
                                                                       
-                 ____  ____  _____    _    _  __
-                | __ )|  _ \| ____|  / \  | |/ /
-                |  _ \| |_) |  _|   / _ \ | ' / 
-                | |_) |  _ <| |___ / ___ \| . \ 
-                |____/|_| \_\_____/_/   \_\_|\_\
-
-      
-Dear users, we are currently adapting the PPIformer Web for the GPU-based predictions.
-Stay tuned for the much faster version of the web application! 
-      
-The new verison of the web will be available withing several days.
-""")
+#                  ____  ____  _____    _    _  __
+#                 | __ )|  _ \| ____|  / \  | |/ /
+#                 |  _ \| |_) |  _|   / _ \ | ' / 
+#                 | |_) |  _ <| |___ / ___ \| . \ 
+#                 |____/|_| \_\_____/_/   \_\_|\_\
+# """)
 
 import sys
 import copy
@@ -45,8 +39,7 @@ 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
+from ppiformer.utils.api import predict_ddg as predict_ddg_
 from ppiformer.utils.torch import fill_diagonal
 from ppiformer.definitions import PPIFORMER_WEIGHTS_DIR
 
@@ -60,14 +53,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):
@@ -288,7 +281,6 @@ 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')