import re, os from pathlib import Path import gradio as gr import spaces import torch 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 device = 'cuda' if torch.cuda.is_available() else 'cpu' @spaces.GPU() def make_uncond_seq(seq_len, model_type): 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=device) 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=device) return generated_sequence def make_cond_seq(seq_len, msa_file, n_sequences, model_type): if model_type == "EvoDiff-MSA": checkpoint = MSA_OA_DM_MAXSUB() model, collater, tokenizer, scheme = checkpoint print(f"MSA File Path: {msa_file.name}") tokeinzed_sample, generated_sequence = generate_query_oadm_msa_simple(msa_file.name, model, tokenizer, int(n_sequences), seq_length=int(seq_len), device=device, selection_type='random') return generated_sequence def make_inpainted_idrs(sequence, start_idx, end_idx, model_type): 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=device) generated_idr_output = { "original_sequence": sequence, "generated_sequence": entire_sequence, "original_region": sequence[start_idx:end_idx], "generated_region": generated_idr } return generated_idr_output # def make_scaffold_motifs(pdb_code, start_idx, end_idx, scaffold_length, model_type): # if model_type == "EvoDiff-Seq": # checkpoint = OA_DM_38M() # model, collater, tokenizer, scheme = checkpoint # data_top_dir = '/home/user/.cache/huggingface/datasets/' # os.makedirs(data_top_dir, exist_ok=True) # # print("Folders in User Cache Directory:", os.listdir("/home/user/.cache")) # 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=device) # generated_scaffold_output = { # "generated_sequence": generated_sequence, # "new_start_index": new_start_idx, # "new_end_index": new_end_idx # } # return generated_scaffold_output usg_app = gr.Interface( fn=make_uncond_seq, inputs=[ gr.Slider(10, 250, step=1, label = "Sequence Length"), gr.Dropdown(["EvoDiff-Seq-OADM 38M", "EvoDiff-D3PM-Uniform 38M"], value="EvoDiff-Seq-OADM 38M", type="value", label = "Model") ], outputs=["text"], 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, 250, label = "Sequence Length"), gr.File(file_types=["a3m"], label = "MSA File"), gr.Number(value=64, precision=0, label = "Number of Sequences to Sample"), gr.Dropdown(["EvoDiff-MSA"], value="EvoDiff-MSA", type="value", label = "Model") ], outputs=["text"], # 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(value = "DQTERTVRSFEGRRTAPYLDSRNVLTIGYGHLLNRPGANKSWEGRLTSALPREFKQRLTELAASQLHETDVRLATARAQALYGSGAYFESVPVSLNDLWFDSVFNLGERKLLNWSGLRTKLESRDWGAAAKDLGRHTFGREPVSRRMAESMRMRRGIDLNHYNI", label = "Sequence"), gr.Number(value=20, precision=0, label = "Start Index"), gr.Number(value=50, precision=0, label = "End Index"), gr.Dropdown(["EvoDiff-Seq"], value="EvoDiff-Seq", type="value", label = "Model") ], outputs=["text"], title = "Inpainting IDRs", description="Inpainting a new region inside a given sequence using the `EvoDiff-Seq` model." ) # scaffold_app = gr.Interface( # fn=make_scaffold_motifs, # inputs=[ # gr.Textbox(value="1prw", label = "PDB Code"), # gr.Textbox(value="[15, 51]", label = "Start Index (as list)"), # gr.Textbox(value="[34, 70]", label = "End Index (as list)"), # gr.Number(value=75, precision=0, label = "Scaffold Length"), # gr.Dropdown(["EvoDiff-Seq", "EvoDiff-MSA"], value="EvoDiff-Seq", type="value", label = "Model") # ], # outputs=["text"], # 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, Ava P. Amini, and Kevin K. Yang] Spaces App By: Tuple, The Cloud Genomics Company [Colby T. Ford] Note: When you first run this app, the models will take a few minutes to download from Zenodo. Check the logs for the download status. """ ) 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()