Update app.py

app.py

@@ -11,7 +11,7 @@ import subprocess
 try:
     import torch_geometric
 except ImportError:
-    subprocess.call(['pip', 'install', 'torch-scatter', 'torch-sparse', 'torch-cluster', 'torch-spline-conv', 'torch-geometric', '-f', 'https://data.pyg.org/whl/torch-1.12.0+cu102.html'])
+    subprocess.call(['pip', 'install', 'torch-scatter', 'torch-sparse', 'torch-cluster', 'torch-spline-conv', 'torch-geometric', '-f', 'https://data.pyg.org/whl/torch-1.12.0+cu113.html'])
 
 import time
 from argparse import ArgumentParser, Namespace, FileType
@@ -243,7 +243,7 @@ def esm(protein_path, out_file):
     )
 
 
-def update(inp, file, ligand_inp, ligand_file, n_it):
+def update(inp, file, ligand_inp, ligand_file, n_it, n_samples, actual_steps, no_final_step_noise):
     pdb_path = get_pdb(inp, file)
     ligand_path = get_ligand(ligand_inp, ligand_file)
 
@@ -270,12 +270,12 @@ def update(inp, file, ligand_inp, ligand_file, n_it):
         [],
         [],
     )
-    N = 10  # number of samples to generate
+    N = n_samples  # number of samples to generate
     protein_path_list = [pdb_path]
     ligand_descriptions = [ligand_path]
     no_random = False
     ode = False
-    no_final_step_noise = False
+    no_final_step_noise = no_final_step_noise
     out_dir = "results/"
     test_dataset = PDBBind(
         transform=None,
@@ -382,7 +382,7 @@ def update(inp, file, ligand_inp, ligand_file, n_it):
             data_list, confidence = sampling(
                 data_list=data_list,
                 model=model,
-                inference_steps=n_it,
+                inference_steps=actual_steps,
                 tr_schedule=tr_schedule,
                 rot_schedule=rot_schedule,
                 tor_schedule=tor_schedule,
@@ -518,9 +518,17 @@ with demo:
                     label="SMILES string",
                 )
                 ligand_file = gr.File(file_count="single", label="Input Ligand")
-        n_it = gr.Slider(
+        n_it = gr.Slider(value=20,
             minimum=10, maximum=40, label="Number of inference steps", step=1
         )
+        actual_steps = gr.Slider(value=18,
+            minimum=10, maximum=40, label="Number of actual inference steps", step=1
+        )
+        n_samples = gr.Slider(value=40,
+            minimum=10, maximum=40, label="Number of samples", step=1
+        )
+        no_final_step_noise = gr.Checkbox(value=True,label="No final step noise"
+        )
 
     btn = gr.Button("Run predictions")
 
@@ -537,52 +545,70 @@ with demo:
                 "examples/6w70.pdb",
                 "COc1ccc(cc1)n2c3c(c(n2)C(=O)N)CCN(C3=O)c4ccc(cc4)N5CCCCC5=O",
                 "examples/6w70_ligand.sdf",
-                10,
+                20,
+                40,
+                18,
+                True
             ],
             [
                 "6moa",
                 "examples/6moa_protein_processed.pdb",
                 "",
                 "examples/6moa_ligand.sdf",
-                10,
+                20,
+                40,
+                18,
+                True
             ],
             [
                 "",
                 "examples/6o5u_protein_processed.pdb",
                 "",
                 "examples/6o5u_ligand.sdf",
-                10,
+                20,
+                40,
+                18,
+                True
             ],
             [
                 "",
                 "examples/6o5u_protein_processed.pdb",
                 "[NH3+]C[C@H]1O[C@H](O[C@@H]2[C@@H]([NH3+])C[C@H]([C@@H]([C@H]2O)O[C@H]2O[C@H](CO)[C@H]([C@@H]([C@H]2O)[NH3+])O)[NH3+])[C@@H]([C@H]([C@@H]1O)O)O",
                 "examples/6o5u_ligand.sdf",
-                10,
+                20,
+                40,
+                18,
+                True
             ],
             [
                 "",
                 "examples/6o5u_protein_processed.pdb",
                 "",
                 "examples/6o5u_ligand.sdf",
-                10,
+                20,
+                40,
+                18,
+                True
             ],
             [
                 "",
                 "examples/6ahs_protein_processed.pdb",
                 "",
                 "examples/6ahs_ligand.sdf",
-                10,
+                20,
+                40,
+                18,
+                True
             ],
         ],
-        [inp, file, ligand_inp, ligand_file, n_it],
+        [inp, file, ligand_inp, ligand_file, n_it, n_samples, actual_steps, no_final_step_noise],
         [mol, out, filenames, pdb, output_file],
         # fn=update,
         # cache_examples=True,
     )
     btn.click(
         fn=update,
-        inputs=[inp, file, ligand_inp, ligand_file, n_it],
+        inputs=[inp, file, ligand_inp, ligand_file, n_it, n_samples, actual_steps, no_final_step_noise],
         outputs=[mol, out, filenames, pdb, output_file],
     )
     out.change(fn=updateView, inputs=[out, filenames, pdb, ligand_file], outputs=mol)