import gradio as gr import os import subprocess from pathlib import Path import shutil import spaces import time import random # Function to check that the files are in the same directory def check_files_in_same_directory(file1, file2): if file1 is None or file2 is None: raise ValueError("Both the YAML config and target PDB/CIF files must be provided.") if Path(file1.name).parent != Path(file2.name).parent: raise ValueError("YAML config and target PDB/CIF must be in the same directory.") def move_file_to_output(file, output_dir): if file is None: return None src_path = Path(file.name) # path to the temp-uploaded file dest_path = Path(output_dir) / src_path.name dest_path.parent.mkdir(parents=True, exist_ok=True) shutil.copy2(src_path, dest_path) # preserves metadata #subprocess.run(["ls", "-l", str(dest_path.parent)], check=True) # Debug: list files in output dir return dest_path def move_files_to_output(yaml_file, target_file, output_dir): yaml_path_dest = move_file_to_output(yaml_file, output_dir) target_path_dest = move_file_to_output(target_file, output_dir) #check_files_in_same_directory(yaml_path_dest, target_path_dest) return yaml_path_dest, target_path_dest def download_results(results_dir, zip_name): """ Creates a ZIP of `results_dir` and returns (zip_path or None, message). """ results_dir = Path(results_dir) if not results_dir.exists() or not results_dir.is_dir(): return None, "No results available for download." # Ensure zip_name is a Path and has .zip extension zip_path = Path(zip_name) if zip_path.suffix != ".zip": zip_path = zip_path.with_suffix(".zip") # Make parent directory if needed if zip_path.parent != Path(""): zip_path.parent.mkdir(parents=True, exist_ok=True) # shutil.make_archive wants base_name without extension base_name = str(zip_path.with_suffix("")) shutil.make_archive( base_name=base_name, format="zip", root_dir=str(results_dir), ) zip_path = Path(base_name + ".zip") return str(zip_path), "Download ready!" def find_log_file(output_dir): # log files start with "boltzgen" and end with ".log" log_files = list(Path(output_dir).glob("boltzgen*.log")) if log_files: # return the most recently modified log file return max(log_files, key=lambda f: f.stat().st_mtime) return None def get_duration(yaml_file, target_file, protocol, num_designs, budget, max_duration): return max_duration @spaces.GPU(duration=get_duration) def run_boltzgen(yaml_file, target_file, protocol, num_designs, budget, max_duration): if yaml_file is None: return "Please upload a YAML config file.", gr.update() timestamp = time.strftime("%Y%m%d-%H%M%S") session_hash = random.getrandbits(128) output_dir = os.path.join(os.getcwd(), "outputs/boltzgen_run_" + timestamp + "_" + str(session_hash)) os.makedirs(output_dir, exist_ok=True) # Move config and target files to the output directory yaml_path_dest, target_path_dest = move_files_to_output(yaml_file, target_file, output_dir) status_message = f"Running BoltzGen with protocol '{protocol}' for {num_designs} designs and budget {budget} tokens." cmd = [ "boltzgen", "run", str(yaml_path_dest), "--output", output_dir, "--protocol", protocol, "--num_designs", str(num_designs), "--budget", str(budget), "--cache", "/tmp/cache" ] print(f"Running BoltzGen with command: {' '.join(cmd)}") status_message += f"\n running command: {' '.join(cmd)}" try: result = subprocess.run(cmd, capture_output=True, text=True, cwd=output_dir) if result.returncode == 0: final_designs_dir = Path(output_dir) / "final_ranked_designs" if final_designs_dir.exists() and final_designs_dir.is_dir(): status_message += "\nBoltzGen completed successfully!" zip_file, _ = download_results(final_designs_dir, "final_ranked_designs.zip") return status_message, zip_file else: status_message += "\nBoltzGen completed but no final designs passed the filtering criteria." return status_message, gr.update() return f"Error:\n{result.stderr}", gr.update() except Exception as e: status_message += f"\nException occurred: {str(e)}" return status_message, gr.update() # Interface with gr.Blocks() as demo: # Markdown with image at the beginning gr.Markdown("""

BoltzGen Design Space

""") gr.Markdown( """ **Design novel protein, peptide, and nanobody binders** against any biomolecular target (proteins, DNA, small molecules) using the BoltzGen diffusion model. **BoltzGen** unifies structure prediction and binder design in a single all-atom generative model with flexible design constraints. Developed by Hannes Stärk and team at MIT: [description](https://jclinic.mit.edu/boltzgen/) [git](https://github.com/HannesStark/boltzgen) """) gr.Markdown( """ ### Instructions 1. Prepare a YAML config file specifying your design task. See [BoltzGen Docs](https://github.com/HannesStark/boltzgen/tree/main/example) for examples. 2. Upload the target structure file (PDB or CIF) corresponding to your design as needed (by the YAML config). 3. Select the design protocol (chose "protein-anything" for general results). 4. Specify the number of designs to generate. !Warning: high numbers may take longer to compute than your session allows. 5. Specify your budget (# of designs wanted in the final diversified set). Note that the budget must be less than or equal to the number of designs. 6. Click "Run BoltzGen" and wait for the results! 7. Download your results using the buttons at the bottom: - Final Ranked Designs: structures, metrics and report summary of the best designs after filtering and ranking (interaction score and diversity). - Intermediate Designs: all generated trajectories before filtering/ranking. """ ) with gr.Row(): yaml_file = gr.File( label="Design YAML Config", file_types=[".yaml", ".yml"], file_count="single" ) target_file = gr.File( label="Target PDB/CIF", file_types=[".pdb", ".cif"], file_count="single" ) with gr.Row(): num_designs = gr.Number( value=10, label="Number of Designs", precision=0, minimum=1, maximum=1000 ) budget = gr.Number( value=5, label="Budget (in tokens)", precision=0, minimum=1, maximum=1000 ) max_duration = gr.Number( value=300, label="Max Duration (seconds)", precision=0, minimum=1, maximum=1800 ) protocol = gr.Dropdown( choices=["protein-anything", "peptide-anything", "nanobody-anything"], value="protein-anything", label="Protocol" ) run_button = gr.Button("Run BoltzGen", variant="primary") output_text = gr.Textbox(label="Run Status", lines=2) download_files = gr.File(label="Download Final Ranked Designs", visible=True) run_button.click( run_boltzgen, inputs=[yaml_file, target_file, protocol, num_designs, budget, max_duration], outputs=[output_text, download_files] ) #gr.Markdown("## Download the results") #with gr.Row(): # download_final_ranked_designs_button = gr.Button("Download Final Results") # download_intermediate_results_button = gr.Button("Download Intermediate Designs") # #with gr.Row(): # download_final_ranked_designs_button.click( # download_results, # inputs=[ # gr.State("/tmp/output/final_ranked_designs"), # gr.State("final_ranked_designs.zip") # ], # outputs=[ # gr.File(label="Download Final Ranked Designs") # ] # ) # # download_intermediate_results_button.click( # download_results, # inputs=[ # gr.State("/tmp/output"), # gr.State("intermediate_results.zip") # ], # outputs=[ # gr.File(label="Download Intermediate Designs") # ] # ) # ## Show the log output in a textbox for debugging #with gr.Row(): # log_output = gr.Textbox(label="BoltzGen Log Output", lines=10) # run_button.click( # lambda: find_log_file("/tmp").read_text() if find_log_file("/tmp") else "No log file found.", # inputs=[], # outputs=log_output # ) demo.queue() demo.launch()