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 fold_all_jobs
#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:



    with gr.Row():
        
        with gr.Column(scale=1.9):  # Left half
            gr.Markdown("# Rosettafold3 for Protein Structure Prediction 🔮 ")
            gr.Markdown("""
            Protein structure prediction has been a fundamental problem in computational biology for decades. <a href="https://www.nature.com/articles/s41586-021-03819-2" target="_blank">Alphafold2's</a> breakthrough in accurately predicting protein structures from sequence alone has revolutionized the field. <a href="https://www.science.org/doi/10.1126/science.abj8754" target="_blank">RoseTTafold</a> provided an open-source alternative with competitive performance. 
            <a href="https://www.nature.com/articles/s41586-024-07487-w" target="_blank">Alphafold3</a> and <a href="https://www.biorxiv.org/content/10.1101/2025.08.14.670328v1" target="_blank">RoseTTafold3</a> build on their predecessors and extend to biomolecular complexes including DNA, RNA and small molecules.
            Structure prediction models reveal unprecedented insights into the function of proteins whose structure has not been resolved experimentally. It is also used in protein design pipelines for in silico screening.  
            This Space pairs well with the companion <a href="https://huggingface.co/spaces/hugging-science/RFDiffusion3" target="_blank">RFDiffusion3</a>, <a href="https://huggingface.co/spaces/hugging-science/LigandMPNN" target="_blank">LigandMPNN</a> and <a href="https://huggingface.co/spaces/hugging-science/ESM2" target="_blank">ESM-2</a> Spaces for a full design pipeline!
                    
            This space allows you to fold bimolecular complexes using Hugging Face's hardware and download the results! It is based on <a href="https://github.com/RosettaCommons/foundry" target="_blank">RosettaCommons' foundry repository.</a>
                        
            <u>Image and Model Source</u>:Corley N, Mathis S, et al. Accelerating Biomolecular Modeling with AtomWorks and RF3. bioRxiv [Preprint]. 2025 Aug 15:2025.08.14.670328. doi: 10.1101/2025.08.14.670328. PMID: 40832246; PMCID: PMC12363939.
            """)

        with gr.Column(scale=1):  # Right half
            gr.Image("assets/rf3.png")

    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 RF3 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.
                

                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 size of the job manageable.
                The time taken significantly depends on the configuration run. An easy way to reduce it is to reduce the number of diffusion steps, which is set to 200 by default but can be reduced to 100 or 50 without a significant drop in performance.
                One can also reduce the number of predictions. These are the only command-line parameters available on this interface, along with the early stopping pLDDT threshold. This threshold stops the model if it becomes obvious during the process that the design is not folding properly.
                It helps save time and resources for erroneous inputs. Its default value of 0.5 is adequate for most use cases. For more information on RosettaFold3's command-line parameters and capabilities, please refer to <a href="https://rosettacommons.github.io/foundry/models/rf3/index.html" target="_blank"> the documentation.</a>

                In this Space, we make the difference between **job files** and **support files**. 
                - **Job files** are directly fed to the model as inputs. They are typically .json files with detailed descriptions of the folding task, but can also be .pdb or .cif files. In these last two cases,
                the model folds all sequences in the file without additional conditioning. It is especially useful for users using this Space for verification of LigandMPNN's designed sequences, 
                as its `backbones` output directory can be directly fed into the space.
                - **Support files** are needed when conditioning a job (e.g templating) or using a multiple sequence alignment (MSA), they are mentioned in the .json job files. They can be .pdb or .cif files containing a ground truth structure or a .a3m.gz file containing the MSA.
                
                Example files can be found in the <a href="https://huggingface.co/spaces/hugging-science/RosettaFold3/tree/main/inputs" target="_blank">`inputs` directory of this repo</a>. They are 
                all taken from <a href="https://github.com/RosettaCommons/foundry/tree/production/models/rf3" target="_blank">RosettaCommons' foundry repository.</a> 
                """)

    # inputs from user
    with gr.Row():
        with gr.Column(scale=1):  # Left half
            job_upload = gr.File(label="Upload Job Files: .json, .pdb, .cif", file_types=[".pdb", ".cif", ".json"], file_count="multiple")

        with gr.Column(scale=1):  # Right half
            support_upload = gr.File(label="Upload Support Files: .pdb, .cif, .gz, .zip", file_types=[".pdb", ".cif", ".gz"], file_count="multiple")

    with gr.Row():
            num_predictions = gr.Number(
                value=5,
                label="Number of Structure Predictions",
                precision=0,
                minimum=1,
                maximum=100
            )
            early_stopping= gr.Number(
                value=0.5,
                label="Early Stopping pLDDT Threshold (0 to disable)",
                precision=2,
                minimum=0,
                maximum=1
            )
            max_duration = gr.Number(
                value=300,
                label="Max Duration (seconds)",
                precision=0,
                minimum=1,
                maximum=3600
            )


    with gr.Accordion(label="Advanced Options", open=False):
        diffusion_steps = gr.Number(
            value=200,
            label="Number of Diffusion Steps",
            precision=0,
            minimum=10,
            maximum=400
        )

    run_btn = gr.Button("Fold", variant="primary")
    runtextbox = gr.Textbox(label="Run status", value="Waiting for folding run...")
    output_file = gr.File(label="Download RF3 results as zip", visible=True)

        
    run_btn.click(
        fold_all_jobs, inputs=[job_upload, support_upload, num_predictions, early_stopping, diffusion_steps, max_duration], outputs=[runtextbox, output_file]
    )


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