Add spaces decorators

app.py

@@ -5,6 +5,7 @@ import os, tempfile
 import torch
 import py3Dmol
 from huggingface_hub import login
+import spaces
 
 from esm.utils.structure.protein_chain import ProteinChain
 from esm.models.esm3 import ESM3
@@ -77,6 +78,7 @@ def render_pdb(pdb_id, chain_id, res_start, res_end, pdb_string=None):
     return Molecule3D(tmp_pdb.name, reps=make_reps(res_start=res_start, res_end=res_end))
 
 ## Function for Scaffolding
+@spaces.GPU()
 def scaffold(model_name, token, pdb_id, chain_id, motif_start, motif_end, prompt_length, insert_size):
     pdb = get_pdb(pdb_id, chain_id)
 
@@ -137,6 +139,7 @@ def scaffold(model_name, token, pdb_id, chain_id, motif_start, motif_end, prompt
     ]
 
 ## Function for Secondary Structure Editing
+@spaces.GPU()
 def ss_edit(model_name, token, pdb_id, chain_id, region_start, region_end, shortened_region_length, shortening_ss8):
     pdb = get_pdb(pdb_id, chain_id)
     edit_region = np.arange(region_start, region_end)
@@ -183,6 +186,7 @@ def ss_edit(model_name, token, pdb_id, chain_id, region_start, region_end, short
         ]
 
 ## Function for SASA Editing
+@spaces.GPU()
 def sasa_edit(model_name, token, pdb_id, chain_id, span_start, span_end, n_samples):
     pdb = get_pdb(pdb_id, chain_id)