prepare for maybe nim speedup with fallback

Author: lucidrains

alphafold3_pytorch/app.py

@@ -3,20 +3,19 @@
 
 import secrets
 import shutil
-import gradio as gr
-from gradio_molecule3d import Molecule3D
 from Bio.PDB import PDBIO
 
 from alphafold3_pytorch import Alphafold3, Alphafold3Input
 
 # constants
+
 model = None
 cache_path = None
 pdb_writer = PDBIO()
 
-
 # main fold function
-def fold(entities, request: gr.Request):
+
+def fold(entities, request):
     proteins = []
     rnas = []
     dnas = []
@@ -59,9 +58,9 @@ def fold(entities, request: gr.Request):
 
     return str(output_path)
 
-
 # gradio
-def delete_cache(request: gr.Request):
+
+def delete_cache(request):
     if not request.session_hash:
         return
 
@@ -71,6 +70,9 @@ def delete_cache(request: gr.Request):
 
 
 def start_gradio_app():
+    import gradio as gr
+    from gradio_molecule3d import Molecule3D
+
     with gr.Blocks(delete_cache=(600, 3600)) as gradio_app:
         entities = gr.State([])
 

alphafold3_pytorch/tensor_typing.py

@@ -1,6 +1,10 @@
 from __future__ import annotations
+
+import sh
 from functools import partial
 import importlib.metadata
+from packaging import version
+
 import torch
 import numpy as np
 
@@ -87,6 +91,18 @@ def package_available(package_name: str) -> bool:
 else:
     checkpoint = partial(torch.utils.checkpoint.checkpoint, use_reentrant = False)
 
+# whether to use Nim or not, depending on if available and version is adequate
+
+try:
+    sh.which('nim')
+    HAS_NIM = True
+    NIM_VERSION = sh.nim(eval = 'echo NimVersion', hints = 'off')
+except sh.ErrorReturnCode_1:
+    HAS_NIM = False
+    NIM_VERSION = None
+
+assert not HAS_NIM or version.parse(NIM_VERSION) >= version.parse('2.0.8'), 'nim version must be 2.0.8 or above'
+
 # check is github ci
 
 IS_GITHUB_CI = env.bool('IS_GITHUB_CI', False)