import re from pathlib import Path 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 from colabfold.batch import run def a3m_file(file): return "tmp.a3m" def predict_protein(sequence): download_alphafold_params("alphafold2_ptm", Path(".")) results = run( queries=[('evodiff_protein', sequence, None)], result_dir='evodiff_protein', use_templates=False, num_relax=0, msa_mode="mmseqs2_uniref_env", model_type="alphafold2_ptm", num_models=1, num_recycles=1, model_order=[1], is_complex=False, data_dir=Path("."), keep_existing_results=False, rank_by="auto", stop_at_score=float(100), zip_results=False, user_agent="colabfold/google-colab-main" ) return f"evodiff_protein/evodiff_protein_unrelaxed_rank_001_alphafold2_ptm_model_1_seed_000.pdb" def display_pdb(path_to_pdb): ''' #function to display pdb in py3dmol SOURCE: https://huggingface.co/spaces/merle/PROTEIN_GENERATOR/blob/main/app.py ''' pdb = open(path_to_pdb, "r").read() view = py3Dmol.view(width=500, height=500) view.addModel(pdb, "pdb") view.setStyle({'model': -1}, {"cartoon": {'colorscheme':{'prop':'b','gradient':'roygb','min':0,'max':1}}})#'linear', 'min': 0, 'max': 1, 'colors': ["#ff9ef0","#a903fc",]}}}) view.zoomTo() output = view._make_html().replace("'", '"') print(view._make_html()) x = f""" {output} """ # do not use ' in this input return f"""""" ''' return f"""""" ''' def make_uncond_seq(seq_len, model_type, pred_structure): if model_type == "EvoDiff-Seq-OADM 38M": checkpoint = OA_DM_38M() model, collater, tokenizer, scheme = checkpoint tokeinzed_sample, generated_sequence = generate_oaardm(model, tokenizer, int(seq_len), batch_size=1, device='cpu') if model_type == "EvoDiff-D3PM-Uniform 38M": checkpoint = D3PM_UNIFORM_38M(return_all=True) model, collater, tokenizer, scheme, timestep, Q_bar, Q = checkpoint tokeinzed_sample, generated_sequence = generate_d3pm(model, tokenizer, Q, Q_bar, timestep, int(seq_len), batch_size=1, device='cpu') if pred_structure: path_to_pdb = predict_protein(generated_sequence) molhtml = display_pdb(path_to_pdb) return generated_sequence, molhtml else: return generated_sequence, None def make_cond_seq(seq_len, msa_file, n_sequences, model_type, pred_structure): if model_type == "EvoDiff-MSA": checkpoint = MSA_OA_DM_MAXSUB() model, collater, tokenizer, scheme = checkpoint tokeinzed_sample, generated_sequence = generate_query_oadm_msa_simple(msa_file.name, model, tokenizer, int(n_sequences), seq_length=int(seq_len), device='cpu', selection_type='random') if pred_structure: path_to_pdb = predict_protein(generated_sequence) molhtml = display_pdb(path_to_pdb) 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, int(start_idx), int(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 = './' start_idx = list(map(int, start_idx.strip('][').split(', '))) end_idx = list(map(int, end_idx.strip('][').split(', '))) 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, inputs=[ gr.Slider(10, 100, step=1, label = "Sequence Length"), gr.Dropdown(["EvoDiff-Seq-OADM 38M", "EvoDiff-D3PM-Uniform 38M"], value="EvoDiff-Seq-OADM 38M", type="value", label = "Model"), gr.Checkbox(value=False, label = "Predict Structure?", visible=False) ], outputs=[ "text", gr.HTML() ], title = "Unconditional sequence generation", description="Generate a sequence with `EvoDiff-Seq-OADM 38M` (smaller/faster) or `EvoDiff-D3PM-Uniform 38M` (larger/slower) models." ) csg_app = gr.Interface( fn=make_cond_seq, inputs=[ gr.Slider(10, 100, label = "Sequence Length"), gr.File(file_types=["a3m"], label = "MSA File"), gr.Number(value=1, placeholder=1, precision=0, label = "Number of Sequences") gr.Dropdown(["EvoDiff-MSA"], value="EvoDiff-MSA", type="value", label = "Model"), gr.Checkbox(value=False, label = "Predict Structure?", visible=False) ], outputs=[ "text", gr.HTML() ], # examples=[["https://github.com/microsoft/evodiff/raw/main/examples/example_files/bfd_uniclust_hits.a3m"]], title = "Conditional sequence generation", 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, precision=0, label = "Start Index"), gr.Number(value=50, placeholder=50, precision=0, 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, precision=0, 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(): gr.Markdown( """ # EvoDiff ## Generation of protein sequences and evolutionary alignments via discrete diffusion models 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: Tuple, The Cloud Genomics Company [Colby T. Ford] """ ) with gr.Row(): gr.TabbedInterface([usg_app, csg_app, idr_app, scaffold_app], ["Unconditional sequence generation", "Conditional generation", "Inpainting IDRs", "Scaffolding functional motifs"]) if __name__ == "__main__": edapp.launch()