diff --git a/README.md b/README.md index 62af505..37797a7 100644 --- a/README.md +++ b/README.md @@ -42,6 +42,12 @@ Boltz currently accepts three input formats: To see all available options: `boltz predict --help` and for more informaton on these input formats, see our [prediction instructions](docs/prediction.md). +### Build easily with Gradio + +``` +python gradio_app.py +``` + ## Training If you're interested in retraining the model, see our [training instructions](docs/training.md). diff --git a/gradio_app.py b/gradio_app.py new file mode 100644 index 0000000..16bdaf3 --- /dev/null +++ b/gradio_app.py @@ -0,0 +1,139 @@ +import spaces +import gradio as gr +from gradio_molecule3d import Molecule3D +from gradio_cofoldinginput import CofoldingInput +import os +import re +import urllib.request +import yaml + +# make sure to pip install gradio gradio_molecule3d gradio_cofoldinginput + +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" + +cache = "/home/user/.boltz" + +os.makedirs(cache) + +ccd = f"{cache}/ccd.pkl" +if not os.path.exists(ccd): + print( + f"Downloading the CCD dictionary to {ccd}. You may " + ) + urllib.request.urlretrieve(CCD_URL, str(ccd)) + +# Download model +model =f"{cache}/boltz1.ckpt" +if not os.path.exists(model): + print( + f"Downloading the model weights to {model}" + ) + urllib.request.urlretrieve(MODEL_URL, str(model)) + + + +@spaces.GPU(duration=120) +def predict(jobname, inputs, recycling_steps, sampling_steps, diffusion_samples): + try: + jobname = re.sub(r'[<>:"/\\|?*]', '_', jobname) + if jobname == "": + raise gr.Error("Job name empty or only invalid characters. Choose a plaintext name.") + os.makedirs(jobname, exist_ok=True) + """format Gradio Component: + # {"chains": [ + # { + # "class": "DNA", + # "sequence": "ATGCGT", + # "chain": "A" + # } + # ], "covMods":[] + # } + """ + #sequences_for_msa = [] + output = { + "sequences": [] + } + representations = [] + 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"]) + if chain["msa"] == False: + sequence_data[entity_type]["msa"] = f"empty" + representations.append({"model":0, "chain":chain["chain"], "style":"cartoon"}) + if entity_type == "ligand": + if "sdf" in chain.keys(): + if chain["sdf"]!="" and chain["name"]=="": + raise gr.Error("Sorry, no SDF support yet.") + if "name" in chain.keys() and len(chain["name"])==3: + sequence_data[entity_type]["ccd"] = chain["name"] + elif "smiles" in chain.keys(): + sequence_data[entity_type]["smiles"] = chain["smiles"] + else: + raise gr.Error("No ligand found, or not in the right format. CCD codes have 3 letters") + + + representations.append({"model":0, "chain":chain["chain"], "style":"stick", "color":"greenCarbon"}) + + 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}") + + os.system(f"boltz predict {jobname}/{jobname}.yaml --use_msa_server --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 Molecule3D(f"{jobname}/boltz_results_{jobname}/predictions/{jobname}/{jobname}_model_0.pdb", label="Output", reps=representations) + except Exception as e: + raise gr.Error(f"failed with error:{e}") + +with gr.Blocks() as blocks: + gr.Markdown("# Boltz-1") + gr.Markdown("""Open GUI for running [Boltz-1 model](https://github.com/jwohlwend/boltz/)
+ Key components: + - MMSeqs2 Webserver [Mirdita et al.](https://www.nature.com/articles/s41592-022-01488-1) + - Boltz-1 Model [Wohlwend et al.](https://github.com/jwohlwend/boltz/) + - Gradio Custom Components [Molecule3D](https://huggingface.co/spaces/simonduerr/gradio_molecule3d)/[Cofolding Input](https://huggingface.co/spaces/simonduerr/gradio_cofoldinginput) by myself + - [3dmol.js Rego & Koes](https://academic.oup.com/bioinformatics/article/31/8/1322/213186) + + 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") + out = Molecule3D(label="Output") + with gr.Tab("Settings"): + recycling_steps =gr.Slider(value=3, minimum=0, label="Recycling steps") + sampling_steps = gr.Slider(value=200, minimum=0, label="Sampling steps") + diffusion_samples = gr.Slider(value=1, label="Diffusion samples") + + gr.Examples([ + ["TOP7",{"chains": [{"class": "protein", "msa":True,"sequence": "MGDIQVQVNIDDNGKNFDYTYTVTTESELQKVLNELMDYIKKQGAKRVRISITARTKKEAEKFAAILIKVFAELGYNDINVTFDGDTVTVEGQLEGGSLEHHHHHH","chain": "A"}], "covMods":[]}], + ["ApixacabanBinderSmiles", {"chains": [{"class": "protein", "msa":True,"sequence": "SVKSEYAEAAAVGQEAVAVFNTMKAAFQNGDKEAVAQYLARLASLYTRHEELLNRILEKARREGNKEAVTLMNEFTATFQTGKSIFNAMVAAFKNGDDDSFESYLQALEKVTAKGETLADQIAKAL","chain": "A"}, {"class":"ligand", "smiles":"COc1ccc(cc1)n2c3c(c(n2)C(=O)N)CCN(C3=O)c4ccc(cc4)N5CCCCC5=O", "sdf":"","name":"","chain": "B"}], "covMods":[]}], + ["ApixacabanBinderCCD", {"chains": [{"class": "protein","msa":True,"sequence": "SVKSEYAEAAAVGQEAVAVFNTMKAAFQNGDKEAVAQYLARLASLYTRHEELLNRILEKARREGNKEAVTLMNEFTATFQTGKSIFNAMVAAFKNGDDDSFESYLQALEKVTAKGETLADQIAKAL","chain": "A"}, {"class":"ligand", "name":"GG2", "sdf":"","chain": "B"}], "covMods":[]}] + ], + inputs = [jobname, inp] + ) + + btn = gr.Button("predict") + + btn.click(fn=predict, inputs=[jobname,inp, recycling_steps, sampling_steps, diffusion_samples], outputs=[out], api_name="predict") + +blocks.launch(ssr_mode=False)