first commit

README.md

@@ -0,0 +1,13 @@
+---
+title: RoseTTAfold2
+emoji: 🏢
+colorFrom: indigo
+colorTo: purple
+sdk: gradio
+sdk_version: 3.33.1
+app_file: app.py
+pinned: false
+license: mit
+---
+
+

app.py

@@ -0,0 +1,660 @@
+import os, time, sys
+
+
+if not os.path.isfile("RF2_apr23.pt"):
+    # send param download into background
+    os.system(
+        "(apt-get install aria2; aria2c -q -x 16 https://colabfold.steineggerlab.workers.dev/RF2_apr23.pt) &"
+    )
+
+if not os.path.isdir("RoseTTAFold2"):
+    print("install RoseTTAFold2")
+    os.system("git clone https://github.com/sokrypton/RoseTTAFold2.git")
+    os.system(
+        "cd RoseTTAFold2/SE3Transformer; pip -q install --no-cache-dir -r requirements.txt; pip -q install ."
+    )
+    os.system(
+        "wget https://raw.githubusercontent.com/sokrypton/ColabFold/beta/colabfold/mmseqs/api.py"
+    )
+
+    # install hhsuite
+    print("install hhsuite")
+    os.makedirs("hhsuite", exist_ok=True)
+    os.system(
+        f"curl -fsSL https://github.com/soedinglab/hh-suite/releases/download/v3.3.0/hhsuite-3.3.0-SSE2-Linux.tar.gz | tar xz -C hhsuite/"
+    )
+
+
+if os.path.isfile(f"RF2_apr23.pt.aria2"):
+    print("downloading RoseTTAFold2 params")
+    while os.path.isfile(f"RF2_apr23.pt.aria2"):
+        time.sleep(5)
+
+os.environ["DGLBACKEND"] = "pytorch"
+sys.path.append("RoseTTAFold2/network")
+if "hhsuite" not in os.environ["PATH"]:
+    os.environ["PATH"] += ":hhsuite/bin:hhsuite/scripts"
+
+import matplotlib.pyplot as plt
+import numpy as np
+from parsers import parse_a3m
+from api import run_mmseqs2
+import py3Dmol
+import torch
+from string import ascii_uppercase, ascii_lowercase
+import hashlib, re, os
+import random
+
+from Bio.PDB import *
+
+
+def get_hash(x):
+    return hashlib.sha1(x.encode()).hexdigest()
+
+
+alphabet_list = list(ascii_uppercase + ascii_lowercase)
+from collections import OrderedDict, Counter
+
+import gradio as gr
+
+if not "pred" in dir():
+    from predict import Predictor
+
+    print("compile RoseTTAFold2")
+    model_params = "RF2_apr23.pt"
+    if torch.cuda.is_available():
+        pred = Predictor(model_params, torch.device("cuda:0"))
+    else:
+        print("WARNING: using CPU")
+        pred = Predictor(model_params, torch.device("cpu"))
+
+
+def get_unique_sequences(seq_list):
+    unique_seqs = list(OrderedDict.fromkeys(seq_list))
+    return unique_seqs
+
+
+def get_msa(seq, jobname, cov=50, id=90, max_msa=2048, mode="unpaired_paired"):
+    assert mode in ["unpaired", "paired", "unpaired_paired"]
+    seqs = [seq] if isinstance(seq, str) else seq
+
+    # collapse homooligomeric sequences
+    counts = Counter(seqs)
+    u_seqs = list(counts.keys())
+    u_nums = list(counts.values())
+
+    # expand homooligomeric sequences
+    first_seq = "/".join(sum([[x] * n for x, n in zip(u_seqs, u_nums)], []))
+    msa = [first_seq]
+
+    path = os.path.join(jobname, "msa")
+    os.makedirs(path, exist_ok=True)
+    if mode in ["paired", "unpaired_paired"] and len(u_seqs) > 1:
+        print("getting paired MSA")
+        out_paired = run_mmseqs2(u_seqs, f"{path}/", use_pairing=True)
+        headers, sequences = [], []
+        for a3m_lines in out_paired:
+            n = -1
+            for line in a3m_lines.split("\n"):
+                if len(line) > 0:
+                    if line.startswith(">"):
+                        n += 1
+                        if len(headers) < (n + 1):
+                            headers.append([])
+                            sequences.append([])
+                        headers[n].append(line)
+                    else:
+                        sequences[n].append(line)
+        # filter MSA
+        with open(f"{path}/paired_in.a3m", "w") as handle:
+            for n, sequence in enumerate(sequences):
+                handle.write(f">n{n}\n{''.join(sequence)}\n")
+        os.system(
+            f"hhfilter -i {path}/paired_in.a3m -id {id} -cov {cov} -o {path}/paired_out.a3m"
+        )
+        with open(f"{path}/paired_out.a3m", "r") as handle:
+            for line in handle:
+                if line.startswith(">"):
+                    n = int(line[2:])
+                    xs = sequences[n]
+                    # expand homooligomeric sequences
+                    xs = ["/".join([x] * num) for x, num in zip(xs, u_nums)]
+                    msa.append("/".join(xs))
+
+    if len(msa) < max_msa and (
+        mode in ["unpaired", "unpaired_paired"] or len(u_seqs) == 1
+    ):
+        print("getting unpaired MSA")
+        out = run_mmseqs2(u_seqs, f"{path}/")
+        Ls = [len(seq) for seq in u_seqs]
+        sub_idx = []
+        sub_msa = []
+        sub_msa_num = 0
+        for n, a3m_lines in enumerate(out):
+            sub_msa.append([])
+            with open(f"{path}/in_{n}.a3m", "w") as handle:
+                handle.write(a3m_lines)
+            # filter
+            os.system(
+                f"hhfilter -i {path}/in_{n}.a3m -id {id} -cov {cov} -o {path}/out_{n}.a3m"
+            )
+            with open(f"{path}/out_{n}.a3m", "r") as handle:
+                for line in handle:
+                    if not line.startswith(">"):
+                        xs = ["-" * l for l in Ls]
+                        xs[n] = line.rstrip()
+                        # expand homooligomeric sequences
+                        xs = ["/".join([x] * num) for x, num in zip(xs, u_nums)]
+                        sub_msa[-1].append("/".join(xs))
+                        sub_msa_num += 1
+            sub_idx.append(list(range(len(sub_msa[-1]))))
+
+        while len(msa) < max_msa and sub_msa_num > 0:
+            for n in range(len(sub_idx)):
+                if len(sub_idx[n]) > 0:
+                    msa.append(sub_msa[n][sub_idx[n].pop(0)])
+                    sub_msa_num -= 1
+                if len(msa) == max_msa:
+                    break
+
+    with open(f"{jobname}/msa.a3m", "w") as handle:
+        for n, sequence in enumerate(msa):
+            handle.write(f">n{n}\n{sequence}\n")
+
+
+from Bio.PDB.PDBExceptions import PDBConstructionWarning
+import warnings
+from Bio.PDB import *
+import numpy as np
+
+
+def add_plddt_to_cif(best_plddts, best_plddt, best_seed, jobname):
+    pdb_parser = PDBParser()
+    warnings.filterwarnings("ignore", category=PDBConstructionWarning)
+    structure = pdb_parser.get_structure(
+        "pdb", f"{jobname}/rf2_seed{best_seed}_00_pred.pdb"
+    )
+    io = MMCIFIO()
+    io.set_structure(structure)
+    io.save(f"{jobname}/rf2_seed{best_seed}_00_pred.cif")
+    plddt_cif = f"""#
+loop_
+_ma_qa_metric.id
+_ma_qa_metric.mode
+_ma_qa_metric.name
+_ma_qa_metric.software_group_id
+_ma_qa_metric.type
+1 global pLDDT 1 pLDDT 
+2 local  pLDDT 1 pLDDT 
+#
+_ma_qa_metric_global.metric_id    1
+_ma_qa_metric_global.metric_value {best_plddt:.3f}
+_ma_qa_metric_global.model_id     1
+_ma_qa_metric_global.ordinal_id   1
+#
+loop_
+_ma_qa_metric_local.label_asym_id
+_ma_qa_metric_local.label_comp_id
+_ma_qa_metric_local.label_seq_id
+_ma_qa_metric_local.metric_id
+_ma_qa_metric_local.metric_value
+_ma_qa_metric_local.model_id
+_ma_qa_metric_local.ordinal_id"""
+
+    for chain in structure[0]:
+        for i, residue in enumerate(chain):
+            plddt_cif += f"\n{chain.id} {residue.resname} {residue.id[1]} 2 {best_plddts[i]*100:.2f} 1 {residue.id[1]}"
+    plddt_cif += "\n#"
+    with open(f"{jobname}/rf2_seed{best_seed}_00_pred.cif", "a") as f:
+        f.write(plddt_cif)
+
+
+def predict(
+    sequence,
+    jobname,
+    sym,
+    order,
+    msa_concat_mode,
+    msa_method,
+    pair_mode,
+    collapse_identical,
+    num_recycles,
+    use_mlm,
+    use_dropout,
+    max_msa,
+    random_seed,
+    num_models,
+    mode="web",
+):
+    if not os.path.exists("/home/user/app"):  # crude check if on spaces
+        if len(sequence) > 600:
+            raise gr.Error(
+                f"Your sequence is too long ({len(sequence)}). "
+                "Please use the full version of RoseTTAfold2 directly from GitHub."
+            )
+    random_seed = int(random_seed)
+    num_models = int(num_models)
+    max_msa = int(max_msa)
+    num_recycles = int(num_recycles)
+    order = int(order)
+
+    max_extra_msa = max_msa * 8
+    sequence = re.sub("[^A-Z:]", "", sequence.replace("/", ":").upper())
+    sequence = re.sub(":+", ":", sequence)
+    sequence = re.sub("^[:]+", "", sequence)
+    sequence = re.sub("[:]+$", "", sequence)
+
+    if sym in ["X", "C"]:
+        copies = int(order)
+    elif sym in ["D"]:
+        copies = int(order) * 2
+    else:
+        copies = {"T": 12, "O": 24, "I": 60}[sym]
+        order = ""
+    symm = sym + str(order)
+
+    sequences = sequence.replace(":", "/").split("/")
+    if collapse_identical:
+        u_sequences = get_unique_sequences(sequences)
+    else:
+        u_sequences = sequences
+    sequences = sum([u_sequences] * copies, [])
+    lengths = [len(s) for s in sequences]
+
+    # TODO
+    subcrop = 1000 if sum(lengths) > 1400 else -1
+
+    sequence = "/".join(sequences)
+    jobname = jobname + "_" + symm + "_" + get_hash(sequence)[:5]
+
+    print(f"jobname: {jobname}")
+    print(f"lengths: {lengths}")
+
+    os.makedirs(jobname, exist_ok=True)
+    if msa_method == "mmseqs2":
+        get_msa(u_sequences, jobname, mode=pair_mode, max_msa=max_extra_msa)
+
+    elif msa_method == "single_sequence":
+        u_sequence = "/".join(u_sequences)
+        with open(f"{jobname}/msa.a3m", "w") as a3m:
+            a3m.write(f">{jobname}\n{u_sequence}\n")
+
+    elif msa_method == "custom_a3m":
+        print("upload custom a3m")
+        # msa_dict = files.upload()
+        lines = msa_dict[list(msa_dict.keys())[0]].decode().splitlines()
+        a3m_lines = []
+        for line in lines:
+            line = line.replace("\x00", "")
+            if len(line) > 0 and not line.startswith("#"):
+                a3m_lines.append(line)
+
+        with open(f"{jobname}/msa.a3m", "w") as a3m:
+            a3m.write("\n".join(a3m_lines))
+
+    best_plddt = None
+    best_seed = None
+    for seed in range(int(random_seed), int(random_seed) + int(num_models)):
+        torch.manual_seed(seed)
+        random.seed(seed)
+        np.random.seed(seed)
+        npz = f"{jobname}/rf2_seed{seed}_00.npz"
+        pred.predict(
+            inputs=[f"{jobname}/msa.a3m"],
+            out_prefix=f"{jobname}/rf2_seed{seed}",
+            symm=symm,
+            ffdb=None,  # TODO (templates),
+            n_recycles=num_recycles,
+            msa_mask=0.15 if use_mlm else 0.0,
+            msa_concat_mode=msa_concat_mode,
+            nseqs=max_msa,
+            nseqs_full=max_extra_msa,
+            subcrop=subcrop,
+            is_training=use_dropout,
+        )
+        plddt = np.load(npz)["lddt"].mean()
+        if best_plddt is None or plddt > best_plddt:
+            best_plddt = plddt
+            best_plddts = np.load(npz)["lddt"]
+            best_seed = seed
+
+        if mode == "web":
+            # Mol* only displays AlphaFold plDDT if they are in a cif.
+            pdb_parser = PDBParser()
+            mmcif_parser = MMCIFParser()
+
+            plddt_cif = add_plddt_to_cif(best_plddts, best_plddt, best_seed, jobname)
+
+            return f"{jobname}/rf2_seed{best_seed}_00_pred.cif"
+        else:
+            # for api just return a pdb file
+            return f"{jobname}/rf2_seed{best_seed}_00_pred.pdb"
+
+
+def predict_api(
+    sequence,
+    jobname,
+    sym,
+    order,
+    msa_concat_mode,
+    msa_method,
+    pair_mode,
+    collapse_identical,
+    num_recycles,
+    use_mlm,
+    use_dropout,
+    max_msa,
+    random_seed,
+    num_models,
+):
+    filename = predict(
+        sequence,
+        jobname,
+        sym,
+        order,
+        msa_concat_mode,
+        msa_method,
+        pair_mode,
+        collapse_identical,
+        num_recycles,
+        use_mlm,
+        use_dropout,
+        max_msa,
+        random_seed,
+        num_models,
+        mode="api",
+    )
+    with open(f"{filename}") as fp:
+        return fp.read()
+
+
+def molecule(input_pdb, public_link):
+    print(input_pdb)
+    print(public_link + "/file=" + input_pdb)
+    link = public_link + "/file=" + input_pdb
+    x = (
+        """<!DOCTYPE html>
+<html lang="en">
+  <head>
+    <meta charset="utf-8" />
+    <meta name="viewport" content="width=device-width, user-scalable=no, minimum-scale=1.0, maximum-scale=1.0">
+    <title>PDBe Molstar - Helper functions</title>
+    <!-- Molstar CSS & JS -->
+    <link rel="stylesheet" type="text/css" href="https://www.ebi.ac.uk/pdbe/pdb-component-library/css/pdbe-molstar-light-3.1.0.css">
+    <script type="text/javascript" src="https://www.ebi.ac.uk/pdbe/pdb-component-library/js/pdbe-molstar-plugin-3.1.0.js"></script>
+    <style>
+      * {
+          margin: 0;
+          padding: 0;
+          box-sizing: border-box;
+      }
+      .msp-plugin ::-webkit-scrollbar-thumb {
+          background-color: #474748 !important;
+      }
+      .viewerSection {
+        margin: 120px 0 0 0px;
+      }
+      #myViewer{
+        float:left;
+        width:100%;
+        height: 800px;
+        position:relative;
+      }
+      .btn{
+      
+                font-family: "Open Sans", sans-serif;
+                    display: inline-block;
+                    outline: none;
+                    cursor: pointer;
+                    font-weight: 600;
+                    border-radius: 3px;
+                    padding: 12px 24px;
+                    border: 0;
+                    margin:0 10px;
+                    line-height: 1.15;
+                    font-size: 16px;
+                    text-decoration: none;
+      }
+      .btn-orange{
+        background: #ff5000;
+        color: #fff;
+                    
+      }
+      .btn-gray{
+                    color: #3a4149;
+                    background: #e7ebee;
+       
+      }
+      .btn:hover{
+      transition: all .1s ease;
+                        box-shadow: 0 0 0 0 #fff, 0 0 0 3px #ddd;}
+    .text-center{
+        display: flex;
+        align-items: center;
+        justify-content: center;
+        padding: 20px 0;
+        }
+    .flex{
+    padding: 10px;
+        display: flex;
+        align-items: center;
+        justify-content: center;
+        width:fit-content; 
+        }
+    .flex svg{ 
+        margin-right: 10px;
+        width:16px;
+        height:16px;
+        }
+        .flex a{
+        margin:0 10px;
+        }
+
+    </style>
+  </head>
+  <body>
+  <div class="text-center">
+    <a class="btn btn-orange flex" href=\""""
+        + link
+        + """\" target="_blank"> <svg fill="none" stroke="currentColor" stroke-width="1.5" viewBox="0 0 24 24" xmlns="http://www.w3.org/2000/svg" aria-hidden="true">
+  <path stroke-linecap="round" stroke-linejoin="round" d="M19.5 13.5L12 21m0 0l-7.5-7.5M12 21V3"></path>
+</svg> <span>CIF File</span></a>
+    <a class="btn btn-gray flex" href=\""""
+        + link.replace(".cif", ".pdb")
+        + """\" target="_blank"> <svg fill="none" stroke="currentColor" stroke-width="1.5" viewBox="0 0 24 24" xmlns="http://www.w3.org/2000/svg" aria-hidden="true">
+  <path stroke-linecap="round" stroke-linejoin="round" d="M19.5 13.5L12 21m0 0l-7.5-7.5M12 21V3"></path>
+</svg>  <span>PDB File</span></a>
+    
+  </div>
+    <div class="viewerSection">
+      <!-- Molstar container -->
+      <div id="myViewer"></div>
+      
+    </div>
+    <script>
+      //Create plugin instance
+      var viewerInstance = new PDBeMolstarPlugin();
+  
+      //Set options (Checkout available options list in the documentation)
+      var options = {
+        customData: {
+          url: \""""
+        + link
+        + """\",
+          format: "cif"
+        },
+        alphafoldView: true,
+        bgColor: {r:255, g:255, b:255},
+        //hideCanvasControls: ["selection", "animation", "controlToggle", "controlInfo"]
+      }
+      
+      //Get element from HTML/Template to place the viewer 
+      var viewerContainer = document.getElementById("myViewer");
+  
+      //Call render method to display the 3D view
+      viewerInstance.render(viewerContainer, options);
+      
+    </script>
+  </body>
+</html>"""
+    )
+
+    return f"""<iframe style="width: 100%; height: 1000px" name="result" allow="midi; geolocation; microphone; camera; 
+    display-capture; encrypted-media;" sandbox="allow-modals allow-forms 
+    allow-scripts allow-same-origin allow-popups 
+    allow-top-navigation-by-user-activation allow-downloads" allowfullscreen="" 
+    allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>"""
+
+
+def predict_web(
+    sequence,
+    jobname,
+    sym,
+    order,
+    msa_concat_mode,
+    msa_method,
+    pair_mode,
+    collapse_identical,
+    num_recycles,
+    use_mlm,
+    use_dropout,
+    max_msa,
+    random_seed,
+    num_models,
+):
+    if os.path.exists("/home/user/app"):
+        public_link = "https://simonduerr-rosettafold2.hf.space/"
+    else:
+        public_link = "http://localhost:7860"
+
+    filename = predict(
+        sequence,
+        jobname,
+        sym,
+        order,
+        msa_concat_mode,
+        msa_method,
+        pair_mode,
+        collapse_identical,
+        num_recycles,
+        use_mlm,
+        use_dropout,
+        max_msa,
+        random_seed,
+        num_models,
+        mode="web",
+    )
+
+    return molecule(filename, public_link)
+
+
+with gr.Blocks() as rosettafold:
+    gr.Markdown("# RoseTTAFold2")
+    gr.Markdown(
+        """If using please cite: [manuscript](https://www.biorxiv.org/content/10.1101/2023.05.24.542179v1) 
+         <br> Heavily based on [RoseTTAFold2 ColabFold notebook](https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/RoseTTAFold2.ipynb)"""
+    )
+    with gr.Accordion("How to use in PyMol", open=False):
+        gr.Markdown(
+            """```os.system('wget https://huggingface.co/spaces/simonduerr/rosettafold2/raw/main/rosettafold_pymol.py')
+run rosettafold_pymol.py
+rosettafold2 sequence, jobname, [sym, order, msa_concat_mode, msa_method, pair_mode, collapse_identical, num_recycles, use_mlm, use_dropout, max_msa, random_seed, num_models]
+color_plddt jobname ```
+"""
+        )
+    sequence = gr.Textbox(
+        label="sequence",
+        value="PIAQIHILEGRSDEQKETLIREVSEAISRSLDAPLTSVRVIITEMAKGHFGIGGELASK",
+    )
+    jobname = gr.Textbox(label="jobname", value="test")
+
+    with gr.Accordion("Additional settings", open=False):
+        sym = gr.Textbox(label="sym", value="X")
+        order = gr.Slider(label="order", value=1, step=1, minimum=1, maximum=12)
+        msa_concat_mode = gr.Dropdown(
+            label="msa_concat_mode",
+            value="default",
+            choices=["diag", "repeat", "default"],
+        )
+
+        msa_method = gr.Dropdown(
+            label="msa_method",
+            value="single_sequence",
+            choices=[
+                "mmseqs2",
+                "single_sequence",
+            ],  # dont allow custom a3m for now , "custom_a3m"
+        )
+        pair_mode = gr.Dropdown(
+            label="pair_mode",
+            value="unpaired_paired",
+            choices=["unpaired_paired", "paired", "unpaired"],
+        )
+
+        num_recycles = gr.Dropdown(
+            label="num_recycles", value="6", choices=["0", "1", "3", "6", "12", "24"]
+        )
+
+        use_mlm = gr.Checkbox(label="use_mlm", value=False)
+        use_dropout = gr.Checkbox(label="use_dropout", value=False)
+        collapse_identical = gr.Checkbox(label="collapse_identical", value=False)
+        max_msa = gr.Dropdown(
+            choices=["16", "32", "64", "128", "256", "512"],
+            value="16",
+            label="max_msa",
+        )
+        random_seed = gr.Textbox(label="random_seed", value=0)
+        num_models = gr.Dropdown(
+            label="num_models", value="1", choices=["1", "2", "4", "8", "16", "32"]
+        )
+
+    btn = gr.Button("Run", visible=False)
+    btn_web = gr.Button("Run")
+
+    output_plain = gr.HTML()
+    output = gr.HTML()
+
+    btn.click(
+        fn=predict_api,
+        inputs=[
+            sequence,
+            jobname,
+            sym,
+            order,
+            msa_concat_mode,
+            msa_method,
+            pair_mode,
+            collapse_identical,
+            num_recycles,
+            use_mlm,
+            use_dropout,
+            max_msa,
+            random_seed,
+            num_models,
+        ],
+        outputs=output_plain,
+        api_name="rosettafold2",
+    )
+    btn_web.click(
+        fn=predict_web,
+        inputs=[
+            sequence,
+            jobname,
+            sym,
+            order,
+            msa_concat_mode,
+            msa_method,
+            pair_mode,
+            collapse_identical,
+            num_recycles,
+            use_mlm,
+            use_dropout,
+            max_msa,
+            random_seed,
+            num_models,
+        ],
+        outputs=output,
+    )
+
+
+rosettafold.launch(share=True, debug=True)

packages.txt

@@ -0,0 +1 @@
+aria2

requirements.txt

@@ -0,0 +1,5 @@
+dgl==1.0.2+cu116 
+matplotlib
+numpy
+torch
+-f https://data.dgl.ai/wheels/cu116/repo.html

rosettafold_pymol.py

@@ -0,0 +1,168 @@
+from pymol import cmd
+import requests
+
+
+# from gradio_client import Client
+
+
+def color_plddt(selection="all"):
+    """
+    AUTHOR
+    Jinyuan Sun
+    https://github.com/JinyuanSun/PymolFold/tree/main
+    MIT License
+
+    DESCRIPTION
+    Colors Predicted Structures by pLDDT
+
+    USAGE
+    color_plddt sele
+
+    PARAMETERS
+
+    sele (string)
+    The name of the selection/object to color by pLDDT. Default: all
+    """
+    # Alphafold color scheme for plddt
+    cmd.set_color("high_lddt_c", [0, 0.325490196078431, 0.843137254901961])
+    cmd.set_color(
+        "normal_lddt_c", [0.341176470588235, 0.792156862745098, 0.976470588235294]
+    )
+    cmd.set_color("medium_lddt_c", [1, 0.858823529411765, 0.070588235294118])
+    cmd.set_color("low_lddt_c", [1, 0.494117647058824, 0.270588235294118])
+
+    # test the scale of predicted_lddt (0~1 or 0~100 ) as b-factors
+    cmd.select("test_b_scale", f"b>1 and ({selection})")
+    b_scale = cmd.count_atoms("test_b_scale")
+
+    if b_scale > 0:
+        cmd.select("high_lddt", f"({selection}) and (b >90 or b =90)")
+        cmd.select("normal_lddt", f"({selection}) and ((b <90 and b >70) or (b =70))")
+        cmd.select("medium_lddt", f"({selection}) and ((b <70 and b >50) or (b=50))")
+        cmd.select("low_lddt", f"({selection}) and ((b <50 and b >0 ) or (b=0))")
+    else:
+        cmd.select("high_lddt", f"({selection}) and (b >.90 or b =.90)")
+        cmd.select(
+            "normal_lddt", f"({selection}) and ((b <.90 and b >.70) or (b =.70))"
+        )
+        cmd.select("medium_lddt", f"({selection}) and ((b <.70 and b >.50) or (b=.50))")
+        cmd.select("low_lddt", f"({selection}) and ((b <.50 and b >0 ) or (b=0))")
+
+    cmd.delete("test_b_scale")
+
+    # set color based on plddt values
+    cmd.color("high_lddt_c", "high_lddt")
+    cmd.color("normal_lddt_c", "normal_lddt")
+    cmd.color("medium_lddt_c", "medium_lddt")
+    cmd.color("low_lddt_c", "low_lddt")
+
+    # set background color
+    cmd.bg_color("white")
+
+
+def query_rosettafold2(
+    sequence: str,
+    jobname: str,
+    sym: str = "X",
+    order: int = 1,
+    msa_concat_mode: str = "diag",
+    msa_method: str = "single_sequence",
+    pair_mode: str = "unpaired_paired",
+    collapse_identical: bool = True,
+    num_recycles: int = 0,
+    use_mlm: bool = True,
+    use_dropout: bool = True,
+    max_msa: int = 16,
+    random_seed: int = 0,
+    num_models: int = 0,
+):
+    """
+    AUTHOR
+    Simon Duerr
+    https://twitter.com/simonduerr
+
+
+    DESCRIPTION
+    Predict a structure using rosettafold2
+
+    USAGE
+    rosettafold2 sequence, jobname, [sym, order, msa_concat_mode, msa_method, pair_mode, collapse_identical, num_recycles, use_mlm, use_dropout, max_msa, random_seed, num_models]
+
+    PARAMETERS
+
+    sequence: (string)
+    one letter amino acid codes that you want to predict
+
+    jobname: string
+    name of the pdbfile that will be outputted
+
+    sym: string
+    symmetry Default: X
+
+    order:
+    Default 1,
+
+    msa_concat_mode:
+    MSA concatenation mode Default:"diag" Options: "diag", "repeat", "default"
+
+    msa_method:
+    MSA method Default:"single_sequence" Options: "mmseqs2", "single_sequence"
+
+    pair_mode:
+    Pair mode Default:"unpaired_paired" Options: "unpaired_paired", "paired", "unpaired"
+
+    collapse_identical:
+    Collapse identical sequences Default:True
+
+    num_recycles:
+    Number of recycles Default:0 Options: 0, 1, 3, 6, 12, 24
+
+    use_mlm:
+    Use MLM Default:True
+
+    use_dropout:
+    Use dropout Default:True
+
+    max_msa:
+    Max MSA Default:16
+
+    random_seed:
+    Random seed Default:0
+
+    num_models:
+    Number of models Default:0
+    """
+    response = requests.post(
+        "http://localhost:7860/run/rosettafold2/",
+        json={
+            "data": [
+                sequence,  # str  in 'sequence' Textbox component
+                jobname,  # str  in 'jobname' Textbox component
+                sym,  # str  in 'sym' Textbox component
+                order,  # int | float (numeric value between 1 and 12) in 'order' Slider component
+                "diag",  # str (Option from: ['diag', 'repeat', 'default']) in 'msa_concat_mode' Dropdown component
+                "single_sequence",  # str (Option from: ['mmseqs2', 'single_sequence', 'custom_a3m']) in 'msa_method' Dropdown component
+                "unpaired_paired",  # str (Option from: ['unpaired_paired', 'paired', 'unpaired']) in 'pair_mode' Dropdown component
+                True,  # bool  in 'collapse_identical' Checkbox component
+                0,  # int (Option from: ['0', '1', '3', '6', '12', '24']) in 'num_recycles' Dropdown component
+                True,  # bool  in 'use_mlm' Checkbox component
+                True,  # bool  in 'use_dropout' Checkbox component
+                16,  # int (Option from: ['16', '32', '64', '128', '256', '512']) in 'max_msa' Dropdown component
+                0,  # int  in 'random_seed' Textbox component
+                1,  # int (Option from: ['1', '2', '4', '8', '16', '32']) in 'num_models' Dropdown component
+            ]
+        },
+    ).json()
+    print(response)
+    try:
+        data = response["data"]
+    except KeyError:
+        print(response["error"])
+        return None
+    with open(f"{jobname}.pdb", "w") as out:
+        out.writelines(data)
+    cmd.load(f"{jobname}.pdb")
+
+
+cmd.extend("rosettafold2", query_rosettafold2)
+cmd.extend("color_plddt", color_plddt)