add num sample steps and ability to sample structure on cuda if available, for @Kaihui-Cheng

Author: lucidrains

alphafold3_pytorch/cli.py

@@ -19,12 +19,16 @@
 @click.option('-prot', '--protein', type = str, multiple = True, help = 'protein sequences')
 @click.option('-rna', '--rna', type = str, multiple = True, help = 'single stranded rna sequences')
 @click.option('-dna', '--dna', type = str, multiple = True, help = 'single stranded dna sequences')
+@click.option('-steps', '--num-sample-steps', type = int, help = 'number of sampling steps to take')
+@click.option('-cuda', '--use-cuda', type = bool, help = 'use cuda if available')
 @click.option('-o', '--output', type = str, help = 'output path', default = 'output.cif')
 def cli(
     checkpoint: str,
     protein: list[str],
     rna: list[str],
     dna: list[str],
+    num_sample_steps: int,
+    use_cuda: bool,
     output: str
 ):
 
@@ -39,8 +43,16 @@ def cli(
 
     alphafold3 = Alphafold3.init_and_load(checkpoint_path)
 
+    if use_cuda and torch.cuda.is_available():
+        alphafold3 = alphafold3.cuda()
+
     alphafold3.eval()
-    structure, = alphafold3.forward_with_alphafold3_inputs(alphafold3_input, return_bio_pdb_structures = True)
+
+    structure, = alphafold3.forward_with_alphafold3_inputs(
+        alphafold3_input,
+        return_bio_pdb_structures = True,
+        num_sample_steps = num_sample_steps
+    )
 
     output_path = Path(output)
     output_path.parents[0].mkdir(exist_ok = True, parents = True)

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.5.17"
+version = "0.5.18"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" },