app.py

import gradio as gr
import warnings
import os
import subprocess
from pathlib import Path
import shutil
import spaces
from atomworks.io.utils.visualize import view
from lightning.fabric import seed_everything
from rfd3.engine import RFD3InferenceConfig, RFD3InferenceEngine
from utils.download_weights import download_weights
from utils.pipelines import test_rfd3_from_notebook, unconditional_generation, generation_with_input_config
#from gradio_molecule3d import Molecule3D
from utils.handle_events import *
from utils.handle_files import *

download_weights()

    
# Gradio UI
with gr.Blocks(title="RFD3 Test") as demo:

    gr.Markdown("# RFdiffusion3 for Backbone generation 🧬 ")

    with gr.Row():
        gr.Markdown("""When the Baker lab released the first version of [RFdiffusion](https://www.nature.com/articles/s41586-023-06415-8), they opened up new avenues for protein design. 
                    The model was based on the previous structure prediction architectures, yet employing the diffusion framework typical of image generation models.
                    It showed impressive results in generating protein backbones for motif scaffolding and binder design.
        
                    
                    
                    Now in its third version, [RFD3](https://pubmed.ncbi.nlm.nih.gov/41000976/) can create binders for an extended set of targets, from DNA/RNA to small molecules, and allows advanced conditiong.
                    This space allows you to run backbone generation jobs using Hugging Face's hardware and download the results!
                    
                    
                    <u>Image and Model Source</u>: Butcher J, Krishna R, Mitra R, Brent RI, Li Y, Corley N, Kim PT, Funk J, Mathis S, Salike S, Muraishi A, Eisenach H, Thompson TR, Chen J, Politanska Y, Sehgal E, Coventry B, Zhang O, Qiang B, Didi K, Kazman M, DiMaio F, Baker D. De novo Design of All-atom Biomolecular Interactions with RFdiffusion3. bioRxiv [Preprint]. 2025 Nov 19:2025.09.18.676967. doi: 10.1101/2025.09.18.676967. PMID: 41000976; PMCID: PMC12458353.""")


        gr.Image("assets/overview_rfd3_baker.png", width=600)

    config_ready = gr.State(None) # whether the config is ready for running generation, takes values None, "manual", or "upload"
    scaffold_ready = gr.State(None) # whether the scaffold is ready for running generation, takes values None, "upload", or "no_input"
    gen_directory = gr.State(None) # the directory where generation results are saved, used to trigger the download of results as zip file
    gen_results = gr.State(None) # the results of the generation, which is a list of dicts where each dict contains batch number "batch", design number "design", path to cif file "cif_path", and path to pdb file "pdb_path".

    # inputs from user
    with gr.Row():
        with gr.Column(scale=1):  # Left half
            gr.Markdown("Set up the configuration for your run through a valid yaml file. Choose the number of batches and designs per batch for your run.")
            config_upload = gr.File(label="Config file: .yaml or .json", file_types=[".pdb", ".yaml", ".json"])
            num_designs_per_batch = gr.Number(
                value=2,
                label="Number of Designs per Batch",
                precision=0,
                minimum=1,
                maximum=8
            )
            num_batches= gr.Number(
                value=5,
                label="Number of Batches",
                precision=0,
                minimum=1,
                maximum=10
            )

            config_validation_btn = gr.Button("Validate Config")
            config_textbox = gr.Textbox(label="Configuration status", value ="Waiting for config validation...")

            with gr.Accordion(label="Advanced Options", open=False):
                extra_args = gr.Textbox(
                    label="Additional CLI Arguments",
                    placeholder="e.g., inference_sampler.step_scale=3 inference_sampler.gamma_0=0.2",
                    lines=3,
                    info="Add extra RFD3 CLI arguments here (optional)"
                )

        with gr.Column(scale=1):  # Right half
            gr.Markdown("Upload your target/scaffold structure as a PDB file to condition the generation. Press 'No Scaffold/Target' if you want to run an unconditional generation.")
            scaffold_upload = gr.File(label="Target/Scaffold PDB", file_types=[".pdb"])
            with gr.Row():
                scaffold_validation_btn = gr.Button("Validate Scaffold")
                no_input_btn = gr.Button("No Scaffold/Target")
            scaffold_textbox = gr.Textbox(label="Scaffold/Target status", value ="Waiting for scaffold validation...")
    run_btn = gr.Button("Run Generation", variant="primary")
    runtextbox = gr.Textbox(label="Run status", value="Waiting for generation run...")


    # validate the configuration
    config_validation_btn.click(validate_config_ready, inputs=[config_upload, num_designs_per_batch, num_batches], outputs=[config_textbox, config_ready])
    scaffold_validation_btn.click(validate_scaffold_ready_with_file, inputs=scaffold_upload, outputs=[scaffold_textbox, scaffold_ready])
    no_input_btn.click(lambda: ("No scaffold/target will be used. Ready for unconditional generation!", "no_input"), outputs=[scaffold_textbox, scaffold_ready])

    output_file = gr.File(label="Download RFD3 results as zip", visible=True)

    # Section to inspect PDB of generated structures
    with gr.Row():
        batch_dropdown = gr.Dropdown(
            choices=[], 
            label="Select Batch",
            visible=True
        )
        design_dropdown = gr.Dropdown(
            choices=[], 
            label="Select Design",
            visible=True
        )
        show_pdb_btn = gr.Button("Show PDB content", visible=True)
        
    display_state = gr.Textbox(label="Selected Batch and Design", visible=True)
    display_state.value = "Please Select a Batch and Design number to show sequence"
    
    def generate(config_ready, scaffold_ready, config_upload, scaffold_upload, num_batches, num_designs_per_batch, extra_args):
        if config_ready is None or scaffold_ready is None:
            return None, None
        if config_ready == "upload" and scaffold_ready == "no_input":
            gen_directory, gen_results = generation_with_input_config(config_upload, scaffold_upload, num_batches, num_designs_per_batch, extra_args)
            return gen_directory, gen_results
        if config_ready=="upload" and scaffold_ready=="upload":
            gen_directory, gen_results = generation_with_input_config(config_upload, scaffold_upload, num_batches, num_designs_per_batch, extra_args)
            return gen_directory, gen_results         
        else:
            return None, None
        
    run_btn.click(give_run_status, inputs=[config_ready, scaffold_ready, num_batches, num_designs_per_batch, config_upload, scaffold_upload], outputs=runtextbox).then(
        generate, inputs=[config_ready, scaffold_ready, config_upload, scaffold_upload, num_batches, num_designs_per_batch, extra_args], outputs=[gen_directory, gen_results]
    ).then(
        update_batch_choices,
        inputs=gen_results,
        outputs=batch_dropdown
    ).then(
        download_results_as_zip,
        inputs=gen_directory,
        outputs=output_file
    ).then(
        lambda gen_results: gr.update(value="Generation completed!") if gen_results is not None else gr.update(), 
        inputs=gen_results,
        outputs=runtextbox
    )
        

    batch_dropdown.change(update_designs, inputs=[batch_dropdown, gen_results], outputs=[design_dropdown])
    design_dropdown.change()
    
    show_pdb_btn.click(show_pdb, inputs=[batch_dropdown, design_dropdown, gen_results], outputs=display_state)



    #def load_viewer(batch, design, result):
    #    if batch is None or design is None:
    #        return gr.update()
    #    pdb_data = next(d["pdb"] for d in result if d["batch"] == int(batch) and d["design"] == int(design))
    #    return gr.update(value=pdb_data, visible=True, reps=[{"style": "cartoon"}])  # Customize style
#
    #visualize_btn.click(load_viewer, inputs=[batch_dropdown, design_dropdown, gen_results], outputs=viewer)

if __name__ == "__main__":
    demo.launch()