Update app.py

app.py

@@ -6,8 +6,11 @@ from gradio_molecule3d import Molecule3D
 from gradio_cofoldinginput import CofoldingInput
 
 import os
+import re
 import urllib.request
 
+from msa import run_mmseqs2
+
 CCD_URL = "https://huggingface.co/boltz-community/boltz-1/resolve/main/ccd.pkl"
 MODEL_URL = "https://huggingface.co/boltz-community/boltz-1/resolve/main/boltz1.ckpt"
 
@@ -34,12 +37,76 @@ if not os.path.exists(model):
 
 @spaces.GPU(duration=120)
 def predict(jobname, inputs, recycling_steps, sampling_steps, diffusion_samples):
-
-    os.system("boltz predict ligand.fasta --output_format pdb")
-    return "boltz_results_ligand/predictions/ligand/ligand_model_0.pdb"
+    jobname = re.sub(r'[<>:"/\\|?*]', '_', jobname)
+    os.makedirs(jobname)
+    """format Gradio Component:
+    # {"chains": [
+    #     {
+    #         "class": "DNA",
+    #         "sequence": "ATGCGT",
+    #         "chain": "A"
+    #     }
+    # ], "covMods":[]
+    # }
+    """
+    sequences_for_msa = []
+    for chain in inputs["chains"]:
+        entity_type = chain["class"].lower()
+        sequence_data = {
+            entity_type: {
+                "id": chain["chain"],
+            }
+        }
+        if entity_type in ["protein", "dna", "rna"]:
+            sequence_data[entity_type]["sequence"] = chain["sequence"]
+            if entity_type == "protein":
+                sequences_for_msa.append(chain["sequence"])
+                sequence_data[entity_type]["msa"] = f"{jobname}/msa.a3m"
+        if entity_type == "ligand":
+            if "sdf" in chains.keys():
+                raise gr.Error("Sorry no SDF support yet")
+            if "name" in chains.keys():
+                 sequence_data[entity_type]["ccd"] = chains["name"]
+            if "smiles" in chains.keys():
+                 sequence_data[entity_type]["smiles"] = chains["smiles"]
+
+        if len(inputs["covMods"])>0:
+            raise gr.Error("Sorry, covMods not supported yet. Coming soon. ")
+        output["sequences"].append(sequence_data)
+
+    # Convert the output to YAML
+    yaml_file_path = f"{jobname}/{jobname}.yaml"
+
+    # Write the YAML output to the file
+    with open(yaml_file_path, "w") as file:
+        yaml.dump(output, file, sort_keys=False, default_flow_style=False)
+
+    os.system(f"cat {yaml_file_path}")
+    a3m_lines_mmseqs2 = run_mmseqs2(
+                    sequences_for_msa,
+                    "./",
+                    use_templates=False,
+                )
+    with open(f"{jobname}/msa.a3m", "w+") as fp:
+        fp.writelines(a3m_lines_mmseqs2)
+
+    os.system(f"boltz predict {jobname}/{jobname}.yaml --out_dir {jobname} --recycling_steps {recycling_steps} --sampling_steps {sampling_steps} --diffusion_samples {diffusion_samples} --override --output_format pdb")
+    print(os.listdir(jobname))
+    print(os.listdir(f"{jobname}/boltz_results_{jobname}/predictions/{jobname}/"))
+    return f"{jobname}/boltz_results_{jobname}/predictions/{jobname}/{jobname}_model_0.pdb"
 
 with gr.Blocks() as blocks:
     gr.Markdown("# Boltz-1")
+    gr.Markdown("""Open GUI for running [Boltz-1 model](https://github.com/jwohlwend/boltz/) <br>
+    Key components:
+    - MMSeqs2 Webserver Mirdita et al. 
+    - Boltz-1 Model Wohlwend et al.
+    - Gradio Custom Components Molecule3D/Cofolding Input Dürr S.
+    - 3dmol.js Rego & Koes 
+    
+    Note: This is an alpha: Some things like covalent modifications or using sdf files don't work yet. You can a Docker image of this on your local infrastructure easily using:
+    `docker run -it -p 7860:7860 --platform=linux/amd64 --gpus all registry.hf.space/simonduerr-boltz-1:latest python app.py`
+    """)
     with gr.Tab("Main"):
         jobname = gr.Textbox(label="Jobname")
         inp = CofoldingInput(label="Input")