Update app.py

app.py

@@ -3,6 +3,7 @@ from transformers.models.esm.openfold_utils.protein import to_pdb, Protein as OF
 from transformers.models.esm.openfold_utils.feats import atom14_to_atom37
 from proteins_viz import *
 import gradio as gr
+import spaces
 
 def convert_outputs_to_pdb(outputs):
     final_atom_positions = atom14_to_atom37(outputs["positions"][-1], outputs)
@@ -39,6 +40,7 @@ torch.backends.cuda.matmul.allow_tf32 = True
 
 model.trunk.set_chunk_size(64)
 
+@spaces.GPU(duration=120)
 def fold_protein(test_protein):
     tokenized_input = tokenizer([test_protein], return_tensors="pt", add_special_tokens=False)['input_ids']
     tokenized_input = tokenized_input.cuda()