app

app.py

@@ -19,6 +19,7 @@ import tempfile
 import shutil
 from Bio import PDB
 from gradio_molecule3d import Molecule3D
+import spaces  # Import spaces for ZeroGPU compatibility
 
 EXAMPLE_PATH = './examples/example.json'
 example_json=[{'sequences': [{'proteinChain': {'sequence': 'MAEVIRSSAFWRSFPIFEEFDSETLCELSGIASYRKWSAGTVIFQRGDQGDYMIVVVSGRIKLSLFTPQGRELMLRQHEAGALFGEMALLDGQPRSADATAVTAAEGYVIGKKDFLALITQRPKTAEAVIRFLCAQLRDTTDRLETIALYDLNARVARFFLATLRQIHGSEMPQSANLRLTLSQTDIASILGASRPKVNRAILSLEESGAIKRADGIICCNVGRLLSIADPEEDLEHHHHHHHH', 'count': 2}}, {'dnaSequence': {'sequence': 'CTAGGTAACATTACTCGCG', 'count': 2}}, {'dnaSequence': {'sequence': 'GCGAGTAATGTTAC', 'count': 2}}, {'ligand': {'ligand': 'CCD_PCG', 'count': 2}}], 'name': '7pzb_need_search_msa'}]
@@ -145,7 +146,6 @@ def plot_3d(pred_loader):
 
     return temp_pdb_path, cif_path
 
-
 def parse_json_input(json_data: List[Dict]) -> Dict:
     """Convert Protenix JSON format to UI-friendly structure"""
     components = {
@@ -209,7 +209,8 @@ def create_protenix_json(input_data: Dict) -> List[Dict]:
     }]
 
 
-
+@torch.inference_mode()
+@spaces.GPU(duration=120)  # Specify a duration to avoid timeout
 def predict_structure(input_collector: dict):
         """Handle both input types"""
         os.makedirs("./output", exist_ok=True)
@@ -286,6 +287,7 @@ runner = InferenceRunner(configs)
 add_watermark = gr.Checkbox(label="Add Watermark", value=True)
 add_watermark1 = gr.Checkbox(label="Add Watermark", value=True)
 
+
 with gr.Blocks(title="FoldMark", css=custom_css) as demo:
     with gr.Row():
         # Use a Column to align the logo and title horizontally
@@ -448,6 +450,8 @@ with gr.Blocks(title="FoldMark", css=custom_css) as demo:
             outputs=[view3d, confidence_plot_image, cif_file]
         )
 
+    @torch.inference_mode()
+    @spaces.GPU(duration=60)
     def is_watermarked(file):
         # Generate a unique subdirectory and filename
         unique_id = str(uuid.uuid4())