app.py

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]

            <span style="color:red">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.</span>
            """
        )
    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()