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 *
#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 Protein Backbone generation 🧬 ")

    with gr.Row():
        gr.Markdown("""When the Baker lab released the first version of <a href="https://www.nature.com/articles/s41586-023-06415-8" target="_blank">RFdiffusion</a>, 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, <a href="https://pubmed.ncbi.nlm.nih.gov/41000976/" target="_blank">RFD3</a> 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)

    gr.Markdown("## How to Use this Space")

    gr.Markdown("""
                Hugging Face spaces is a great tool to build quick machine learning demos and expose interfaces for popular models. This space runs the RFD3 model with a convenient interface for input configurations and inference parameters.
                The green `Running on ZERO` at the top of the page indicates that this space runs on Hugging Face's <a href="https://huggingface.co/docs/hub/en/spaces-zerogpu" target="_blank">ZeroGPU technology</a>. This means that all jobs are run on NVIDIA H200 GPUs and that the number of GPUs allocated to this space scale automatically based on demand.
                
                There are three steps to setting up a job.
                - First, create an input specification for your generation job. This is a yaml or json file with details on targeted motifs, scaffolds and generation parameters that gets passed to RFD3's CLI as `inputs`. Check out <a href="https://rosettacommons.github.io/foundry/models/rfd3/input.html" target="_blank">RFD3's documentation and tutorials</a> to learn how to compose this file!
                - Second, upload a PDB file containing your target/scaffold structure to condition the generation or leave empty for unconditional generation.
                - Third, Select the number of batches and designs per batch for your job. These are passed as CLI arguments to RFD3. Advanced CLI arguments can be added in the "Advanced Options" section. While most of these arguments can be kept at their default value, some  related to the inference sampler can provide improvements for protein-protein interaction (PPI) workflows. Here is a list of other <a href="https://rosettacommons.github.io/foundry/models/rfd3/input.html#cli-arguments" target="_blank">CLI arguments</a>.

                
                Individual jobs On ZeroGPU spaces are limited to 240 seconds for PRO users or members of an organization. Hence, make sure that you are logged in before launching a job to avoid errors and keep the number of designs per run manageable.
                While the time taken significantly depends on the configuration run, here are some numbers to guide you:
                - RFD3's PPI tutorial, designing a 190-270 aa binder for a 149 aa motif on the <a href="https://www.rcsb.org/structure/4ZXB" target="_blank">human insulin receptor</a>
                generated a batch of 8 designs every 60 seconds, a batch of 16 every 95 seconds.
                - RFD3's complex nucleic acid tutorial, scaffolding a DNA-binding motif of 8 residues into a 150 aa binder,
                generated a batch of 8 designs every 52 seconds.
                - RDF3's enzyme design tutorial, scaffolding the active site of drosophila's alcohol dehydrogenase into a 180-200 aa enzyme, generated a batch of 8 designs every 60 seconds.
                
                
                Example files for these three tutorials can be found in the <a href="https://huggingface.co/spaces/hugging-science/RFdiffusion3/tree/main/examples" target="_blank">`examples` directory of this repo</a>. They are 
                all taken from <a href="https://github.com/RosettaCommons/foundry/tree/production/models/rfd3" target="_blank">RFD3's amazing Github repository.</a> This space pairs well with the <a href="https://huggingface.co/spaces/hugging-science/LigandMPNN" target="_blank">companion LigandMPNN space</a> for inverse folding!

                """)


    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
            config_upload = gr.File(label="Config file: .yaml or .json", file_types=[".pdb", ".yaml", ".json"])

        with gr.Column(scale=1):  # Right half
            scaffold_upload = gr.File(label="Target/Scaffold PDB", file_types=[".pdb"])

    with gr.Row():
            num_designs_per_batch = gr.Number(
                value=8,
                label="Number of Designs per Batch",
                precision=0,
                minimum=1,
                maximum=16
            )
            num_batches= gr.Number(
                value=2,
                label="Number of Batches",
                precision=0,
                minimum=1,
                maximum=50
            )
            max_duration = gr.Number(
                value=300,
                label="Max Duration (seconds)",
                precision=0,
                minimum=1,
                maximum=3600
            )

    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.Row():
        run_btn = gr.Button("Run Generation", variant="primary")
        #stop_btn = gr.Button("Stop Generation", variant="stop", visible=False)
    runtextbox = gr.Textbox(label="Run status", value="Waiting for generation run...")


    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"
        
    #generation_event = run_btn.click(
    #    lambda: (gr.update(visible=False), gr.update(visible=True)),
    #    outputs=[run_btn, stop_btn]
    #).then(
    generation_event = run_btn.click(
        generation_with_input_config, inputs=[config_upload, scaffold_upload, num_batches, num_designs_per_batch, extra_args, max_duration], outputs=[runtextbox, gen_results, output_file]
    ).then(
        update_batch_choices,
        inputs=gen_results,
        outputs=batch_dropdown
    )#.then(
    #    lambda: (gr.update(visible=True), gr.update(visible=False)),
    #    outputs=[run_btn, stop_btn]
    #)
    
    #stop_btn.click(
    #    lambda: gr.update(value="Generation cancelled by user."),
    #    outputs=runtextbox,
    #    cancels=[generation_event]
    #).then(
    #    lambda: (gr.update(visible=True), gr.update(visible=False)),
    #    outputs=[run_btn, stop_btn]
    #)
        

    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.queue()
    demo.launch()