Update app.py

app.py

@@ -1,6 +1,7 @@
 from transformers import AutoTokenizer, EsmForProteinFolding
 from transformers.models.esm.openfold_utils.protein import to_pdb, Protein as OFProtein
 from transformers.models.esm.openfold_utils.feats import atom14_to_atom37
+from Bio import SeqIO
 import gradio as gr
 import spaces
 from gradio_molecule3d import Molecule3D
@@ -129,8 +130,26 @@ def fold_protein(test_protein):
     html = molecule("output_structure.pdb")
     return html, "output_structure.pdb"
 
+@spaces.GPU(duration=180)
+def fold_protein_wpdb(test_protein, pdb_path):
+    tokenized_input = tokenizer([test_protein], return_tensors="pt", add_special_tokens=False)['input_ids']
+    tokenized_input = tokenized_input.cuda()
+    with torch.no_grad():
+        output = model(tokenized_input)
+    pdb = convert_outputs_to_pdb(output)
+    with open(pdb_path, "w") as f:
+        f.write("".join(pdb))
+    html = molecule(pdb_path)
+    return html, pdb_path
+
+def load_protein_sequences(fasta_file):
+    protein_sequences = {} 
+    for record in SeqIO.parse(fasta_file, "fasta"):
+        protein_sequences[record.id] = str(record.seq)
+    return protein_sequences
+
 iface = gr.Interface(
-    title="everything-ai-proteinfold",
+    title="Proteinviz",
     fn=fold_protein,
     inputs=gr.Textbox(
             label="Protein Sequence",
@@ -146,4 +165,20 @@ iface = gr.Interface(
     ]
 )
 
-iface.launch(server_name="0.0.0.0", share=False)
+with gr.Blocks() as demo1:
+    input_seqs = gr.File(label="FASTA File", info="FASTA-formatted file with headers starting with '>' and protein sequences")
+    @gr.render(inputs=input_text)
+    def show_split(inputfile):
+        if type(inputfile) == type(None):
+            gr.Markdown("## No Input Provided")
+        else:
+            seqs = load_protein_sequences(inputfile)
+            for header in seqs:
+                pdb_path = f"{seq.replace(" ", "_").replace(",","")}.pdb"
+                html, pdb = fold_protein_wpdb(seqs[seq], pdb_path)
+                gr.HTML(html, label=f"{seq} structural representation")   
+                Molecule3D(pdb, label=f"{seq} molecular representation")
+
+demo = gr.TabbedInterface([iface, demo0], ["Single Protein Structure Prediction", "Bulk Protein Structure Prediction"])
+
+demo.launch(server_name="0.0.0.0", share=False)