Update inference_app.py

inference_app.py

@@ -8,10 +8,10 @@ from gradio_molecule3d import Molecule3D
 
 
 
-def predict (input_sequence, input_ligand):
+def predict (input_sequence, input_ligand, input_protein):
     start_time = time.time()
     # Do inference here
-    # return an output directory
+    # return an output pdb file with the protein and ligand with resname LIG or UNK. 
     end_time = time.time()
     run_time = end_time - start_time
     return "test_out.pdb", run_time
@@ -57,12 +57,19 @@ with gr.Blocks() as app:
       "resname": "UNK",
       "style": "stick",
       "color": "greenCarbon",
+    },
+        {
+      "model": 0,
+      "resname": "LIG",
+      "style": "stick",
+      "color": "greenCarbon",
     }
+        
   ]
     
     out = Molecule3D(reps=reps)
     run_time = gr.Textbox(label="Runtime")
 
-    btn.click(predict, inputs=[input_sequence, input_ligand], outputs=[out, run_time])
+    btn.click(predict, inputs=[input_sequence, input_ligand, input_protein], outputs=[out, run_time])
 
 app.launch()