Update app.py

app.py

@@ -5,6 +5,7 @@ import gradio as gr
 from evodiff.pretrained import OA_DM_38M, D3PM_UNIFORM_38M, MSA_OA_DM_MAXSUB
 from evodiff.generate import generate_oaardm, generate_d3pm
 from evodiff.generate_msa import generate_query_oadm_msa_simple
+from evodiff.conditional_generation import inpaint_simple, generate_scaffold
 
 import py3Dmol
 from colabfold.download import download_alphafold_params
@@ -97,6 +98,49 @@ def make_cond_seq(seq_len, msa_file, model_type, pred_structure):
         return generated_sequence, molhtml
     else:
         return generated_sequence, None
+    
+def make_inpainted_idrs(sequence, start_idx, end_idx, model_type, pred_structure):
+    if model_type == "EvoDiff-Seq":
+        checkpoint = OA_DM_38M()
+        model, collater, tokenizer, scheme = checkpoint
+        sample, entire_sequence, generated_idr = inpaint_simple(model, sequence, start_idx, end_idx, tokenizer=tokenizer, device='cpu')
+
+        generated_idr_output = {
+            "original_sequence": sequence,
+            "generated_sequence": entire_sequence,
+            "original_region": sequence[start_idx:end_idx],
+            "generated_region": generated_idr
+        }
+
+    if pred_structure:
+        path_to_pdb = predict_protein(entire_sequence)
+        molhtml = display_pdb(path_to_pdb)
+
+        return generated_idr_output, molhtml
+    else:
+        return generated_idr_output, None
+    
+def make_scaffold_motifs(pdb_code, start_idx, end_idx, scaffold_length, model_type, pred_structure):
+    if model_type == "EvoDiff-Seq":
+        checkpoint = OA_DM_38M()
+        model, collater, tokenizer, scheme = checkpoint
+        data_top_dir = './'
+        generated_sequence, new_start_idx, new_end_idx = generate_scaffold(model, pdb_code, start_idx, end_idx, scaffold_length, data_top_dir, tokenizer, device='cpu')
+
+        generated_scaffold_output = {
+            "generated_sequence": generated_sequence,
+            "new_start_index": new_start_idx,
+            "new_end_index": new_end_idx
+        }
+
+    if pred_structure:
+        # path_to_pdb = predict_protein(generated_sequence)
+        path_to_pdb = f"scaffolding-pdbs/{pdb_code}.pdb"
+        molhtml = display_pdb(path_to_pdb)
+
+        return generated_scaffold_output, molhtml
+    else:
+        return generated_scaffold_output, None
 
 usg_app = gr.Interface(
             fn=make_uncond_seq,
@@ -130,6 +174,40 @@ csg_app = gr.Interface(
             description="Evolutionary guided sequence generation with the `EvoDiff-MSA` model."
             )
 
+idr_app = gr.Interface(
+            fn=make_inpainted_idrs,
+            inputs=[
+                gr.Textbox(placeholder="DQTERTVRSFEGRRTAPYLDSRNVLTIGYGHLLNRPGANKSWEGRLTSALPREFKQRLTELAASQLHETDVRLATARAQALYGSGAYFESVPVSLNDLWFDSVFNLGERKLLNWSGLRTKLESRDWGAAAKDLGRHTFGREPVSRRMAESMRMRRGIDLNHYNI", label = "Sequence"),
+                gr.Number(value=20, placeholder=20, label = "Start Index"),
+                gr.Number(value=50, placeholder=50, label = "End Index"),
+                gr.Dropdown(["EvoDiff-Seq"], value="EvoDiff-Seq", type="value", label = "Model"),
+                gr.Checkbox(value=False, label = "Predict Structure?", visible=False)
+                ],
+            outputs=[
+                "text",
+                gr.HTML()
+            ],
+            title = "Inpainting IDRs",
+            description="Inpaining a new region inside a given sequence using the `EvoDiff-Seq` model."
+            )
+
+scaffold_app = gr.Interface(
+            fn=make_scaffold_motifs,
+            inputs=[
+                gr.Textbox(placeholder="1prw", label = "PDB Code"),
+                gr.Textbox(value="[15, 51]", placeholder="[15, 51]", label = "Start Index (as list)"),
+                gr.Textbox(value="[34, 70]", placeholder="[34, 70]", label = "End Index (as list)"),
+                gr.Number(value=75, placeholder=75, label = "Scaffold Length"),
+                gr.Dropdown(["EvoDiff-Seq", "EvoDiff-MSA"], value="EvoDiff-Seq", type="value", label = "Model"),
+                gr.Checkbox(value=False, label = "Predict Structure?", visible=False)
+                ],
+            outputs=[
+                "text",
+                gr.HTML()
+            ],
+            title = "Scaffolding functional motifs",
+            description="Scaffolding a new functional motif inside a given PDB structure using the `EvoDiff-Seq` model."
+            )
 
 with gr.Blocks() as edapp:
     with gr.Row():
@@ -140,11 +218,15 @@ with gr.Blocks() as edapp:
 
             Created By: Microsoft Research [Sarah Alamdari, Nitya Thakkar, Rianne van den Berg, Alex X. Lu, Nicolo Fusi, ProfileAva P. Amini, and Kevin K. Yang]
             
-            Spaces App By: [Colby T. Ford](httos://github.com/colbyford)
+            Spaces App By: Tuple, The Cloud Genomics Company [Colby T. Ford]
             """
         )
     with gr.Row():
-        gr.TabbedInterface([usg_app, csg_app], ["Unconditional sequence generation", "Conditional generation"])
+        gr.TabbedInterface([usg_app, csg_app, idr_app, scaffold_app],
+                           ["Unconditional sequence generation",
+                            "Conditional generation",
+                            "Inpainting IDRs",
+                            "Scaffolding functional motifs"])