unconditional generation with input file pipeline

app.py

@@ -9,7 +9,7 @@ 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
+from utils.pipelines import test_rfd3_from_notebook, unconditional_generation, unconditional_generation_with_input_config
 #from gradio_molecule3d import Molecule3D
 from utils.handle_events import *
 from utils.handle_files import *
@@ -52,7 +52,7 @@ with gr.Blocks(title="RFD3 Test") as demo:
             gr.Markdown("Set up the configuration for your run through a valid yaml file or by manually setting minimal parameters for an unconditional run.")
             with gr.Tabs() as config_tabs:
                 with gr.TabItem("Upload Config") as upload_tab: # upload a config yaml or json
-                    config_upload = gr.File(label="PDB + Config", file_types=[".pdb", ".yaml", ".json"])
+                    config_upload = gr.File(label="Config file: .yaml or .json", file_types=[".pdb", ".yaml", ".json"])
                 with gr.TabItem("Manual Config") as manual_tab: # minimal config for testing unconditional generation
                     num_designs_per_batch = gr.Number(
                         value=2,
@@ -123,6 +123,9 @@ with gr.Blocks(title="RFD3 Test") as demo:
     def generate(config_ready, scaffold_ready, num_batches, num_designs_per_batch, length):
         if config_ready is None or scaffold_ready is None:
             return None, None
+        if config_ready is "upload" and scaffold_ready is "no_input":
+            gen_directory, gen_results = unconditional_generation_with_input_config(config_upload)
+            return gen_directory, gen_results
         if config_ready=="manual" and scaffold_ready=="no_input":
             gen_directory, gen_results = unconditional_generation(num_batches, num_designs_per_batch, length)
             return gen_directory, gen_results

utils/handle_files.py

@@ -2,6 +2,7 @@ import gemmi
 import os
 import shutil
 import gradio as gr
+import subprocess
 
 def mcif_gz_to_pdb(file_path: str) -> str:
     """
@@ -40,4 +41,13 @@ def download_results_as_zip(directory):
         return gr.update()
     zip_path = f"{directory}.zip"
     shutil.make_archive(directory, 'zip', directory)
-    return gr.update(value=zip_path, visible=True)
+    return gr.update(value=zip_path, visible=True)
+
+
+def collect_outputs(gen_directory, num_batches, num_designs_per_batch):
+    try:
+        cmd = f"ls -R {gen_directory}"
+        file_list = subprocess.check_output(cmd, shell=True).decode()
+        return file_list
+    except Exception as e:
+        return f"Error: {str(e)}"

utils/pipelines.py

@@ -46,6 +46,26 @@ def test_rfd3_from_notebook():
 # Initialize engine and run generation
 @spaces.GPU(duration=300) 
 def unconditional_generation(num_batches, num_designs_per_batch, length):
+    """
+    Runs an unconditional generation with the specified parameters for number of batches, number of designs per batch, and length of the generated proteins. Saves the generated structures to a timestamped directory in the outputs folder and returns the path to the directory along with a list of the generated structures' file paths.
+
+    Parameters:
+    ----------
+    num_batches: int or gr.Number,
+        The number of batches to generate.
+    num_designs_per_batch: int or gr.Number,
+        The number of designs to generate per batch.
+    length: int or gr.Number,
+        The length of the generated proteins.
+
+    Returns:
+    -------
+        directory: str,
+            The path to the directory where the generated structures are saved.
+        results: list of dicts,
+            A list of the generated structures' file paths, where each dict contains batch number "batch", design number "design", path to cif file "cif_path", and path to pdb file "pdb_path".
+    """
+
 
     config = RFD3InferenceConfig(
     specification={
@@ -79,12 +99,50 @@ def unconditional_generation(num_batches, num_designs_per_batch, length):
     except Exception as e:
         raise RuntimeError(f"Error during generation: {str(e)}")
     
-def collect_outputs(gen_directory, num_batches, num_designs_per_batch):
-    try:
-        cmd = f"ls -R {gen_directory}"
-        file_list = subprocess.check_output(cmd, shell=True).decode()
-        return file_list
+def unconditional_generation_with_input_config(input_file):
+    """
+    Runs an unconditional generation with the specified input config file. Saves the generated structures to a timestamped directory in the outputs folder and returns the path to the directory along with a list of the generated structures' file paths.
+
+    Parameters:
+    ----------
+    input_file: gr.File,
+        gr.File object containing the uploaded config file (yaml or json). input_file.name is the path to the uploaded file on the server.
+    
+    Returns:
+    -------
+        directory: str,
+            The path to the directory where the generated structures are saved.
+        results: list of dicts,
+            A list of the generated structures' file paths, where each dict contains batch number "batch", design number "design", path to cif file "cif_path", and path to pdb file "pdb_path".
+
+    """
+
+    session_hash = gr.Request().session_hash
+    time_stamp = time.strftime("%Y-%m-%d-%H-%M-%S")
+    directory = f"./outputs/unconditional_generation/session_{session_hash}_{time_stamp}"
+    os.makedirs(directory, exist_ok=False)
+
+    try: 
+        
+        command = f"rfd3 design inputs={input_file.name} out_dir={directory}"
+        print(f"Running command: {command}")
+        subprocess.run(command, shell=True, check=True)
+
+        results = []
+        for file_name in os.listdir(directory):
+            if file_name.endswith(".cif.gz"):
+                terms = file_name.split("_")
+                model_index = terms.index("model")
+                batch =model_index - 1
+                design = model_index + 1
+                cif_path = os.path.join(directory, file_name)
+                pdb_path = mcif_gz_to_pdb(cif_path)
+                results.append({"batch": batch, "design": design, "cif_path": cif_path, "pdb_path": pdb_path})
+            
+        print(results)
+        return directory, results
+    
     except Exception as e:
-        return f"Error: {str(e)}"
+        raise RuntimeError(f"Error during generation: {str(e)}")