add add_data_args to inference and config relax

openfold/np/relax/amber_minimize.py

@@ -105,7 +105,7 @@ def _openmm_minimize(
         _add_restraints(system, pdb, stiffness, restraint_set, exclude_residues)
 
     integrator = openmm.LangevinIntegrator(0, 0.01, 0.0)
-    platform = openmm.Platform.getPlatformByName("CUDA" if use_gpu else "CPU")
+    platform = openmm.Platform.getPlatformByName("HIP" if use_gpu else "CPU")
     simulation = openmm_app.Simulation(
         pdb.topology, system, integrator, platform
     )

openfold/utils/script_utils.py

@@ -235,12 +235,12 @@ def relax_protein(config, model_device, unrelaxed_protein, output_directory, out
     )
 
     t = time.perf_counter()
-    visible_devices = os.getenv("CUDA_VISIBLE_DEVICES", default="")
+    visible_devices = os.getenv("HIP_VISIBLE_DEVICES", default="")
     if "cuda" in model_device:
         device_no = model_device.split(":")[-1]
-        os.environ["CUDA_VISIBLE_DEVICES"] = device_no
+        os.environ["HIP_VISIBLE_DEVICES"] = device_no
     relaxed_pdb_str, _, _ = amber_relaxer.process(prot=unrelaxed_protein)
-    os.environ["CUDA_VISIBLE_DEVICES"] = visible_devices
+    os.environ["HIP_VISIBLE_DEVICES"] = visible_devices
     relaxation_time = time.perf_counter() - t
 
     logger.info(f"Relaxation time: {relaxation_time}")

run_pretrained_openfold.py

@@ -55,7 +55,48 @@ from openfold.utils.trace_utils import (
     pad_feature_dict_seq,
     trace_model_,
 )
-from scripts.utils import add_data_args
+from datetime import date
+
+# from scripts.utils import add_data_args
+
+def add_data_args(parser: argparse.ArgumentParser):
+    parser.add_argument(
+        '--uniref90_database_path', type=str, default=None,
+    )
+    parser.add_argument(
+        '--mgnify_database_path', type=str, default=None,
+    )
+    parser.add_argument(
+        '--pdb70_database_path', type=str, default=None,
+    )
+    parser.add_argument(
+        '--uniclust30_database_path', type=str, default=None,
+    )
+    parser.add_argument(
+        '--bfd_database_path', type=str, default=None,
+    )
+    parser.add_argument(
+        '--jackhmmer_binary_path', type=str, default='/usr/bin/jackhmmer'
+    )
+    parser.add_argument(
+        '--hhblits_binary_path', type=str, default='/usr/bin/hhblits'
+    )
+    parser.add_argument(
+        '--hhsearch_binary_path', type=str, default='/usr/bin/hhsearch'
+    )
+    parser.add_argument(
+        '--kalign_binary_path', type=str, default='/usr/bin/kalign'
+    )
+    parser.add_argument(
+        '--max_template_date', type=str,
+        default=date.today().strftime("%Y-%m-%d"),
+    )
+    parser.add_argument(
+        '--obsolete_pdbs_path', type=str, default=None
+    )
+    parser.add_argument(
+        '--release_dates_path', type=str, default=None
+    )
 
 
 TRACING_INTERVAL = 50

run_pretrained_openfold.sh

+python3 run_pretrained_openfold.py \
+    fasta \
+    /data/pdb_mmcif/mmcif_files/ \
+    --uniref90_database_path /data/uniref90/uniref90.fasta \
+    --mgnify_database_path /data/mgnify/mgy_clusters_2018_12.fa \
+    --pdb70_database_path /data/pdb70/pdb70 \
+    --uniclust30_database_path /data/uniclust30/uniclust30_2018_08/uniclust30_2018_08 \
+    --output_dir ./ \
+    --bfd_database_path /data/bfd/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt \
+    --model_device "cuda:0" \
+    --jackhmmer_binary_path `which jackhmmer` \
+    --hhblits_binary_path `which hhblits` \
+    --hhsearch_binary_path `which hhsearch` \
+    --kalign_binary_path `which kalign` \
+    --config_preset "model_1" \
+    --jax_param_path /data/params/params_model_1.npz \
+    --use_precomputed_alignments alignments/