Update inference_app.py

inference_app.py

@@ -8,13 +8,13 @@ from gradio_molecule3d import Molecule3D
 
 
 
-def predict (input_sequence, input_ligand, input_protein):
+def predict (input_sequence, input_ligand,input_msa, input_protein):
     start_time = time.time()
     # Do inference here
     # 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
+    return ["test_out.pdb", "test_docking_pose.sdf"], run_time
 
 with gr.Blocks() as app:
 
@@ -24,6 +24,8 @@ with gr.Blocks() as app:
     with gr.Row():
         input_sequence = gr.Textbox(lines=3, label="Input Protein sequence (FASTA)")
         input_ligand = gr.Textbox(lines=3, label="Input ligand SMILES")
+    with gr.Row():
+        input_msa = gr.File(label="Input Protein MSA (FASTA)")
         input_protein = gr.File(label="Input protein monomer")
         
     
@@ -51,16 +53,9 @@ with gr.Blocks() as app:
       "model": 0,
       "style": "cartoon",
       "color": "whiteCarbon",
-    },
-    {
-      "model": 0,
-      "resname": "UNK",
-      "style": "stick",
-      "color": "greenCarbon",
     },
         {
-      "model": 0,
-      "resname": "LIG",
+      "model": 1,
       "style": "stick",
       "color": "greenCarbon",
     }
@@ -70,6 +65,6 @@ with gr.Blocks() as app:
     out = Molecule3D(reps=reps)
     run_time = gr.Textbox(label="Runtime")
 
-    btn.click(predict, inputs=[input_sequence, input_ligand, input_protein], outputs=[out, run_time])
+    btn.click(predict, inputs=[input_sequence, input_ligand, input_msa, input_protein], outputs=[out, run_time])
 
 app.launch()