move generation scripts to pipelines, only download rfd3, create output directory for run

app.py

@@ -9,44 +9,11 @@ from atomworks.io.utils.visualize import view
 from lightning.fabric import seed_everything
 from rfd3.engine import RFD3InferenceConfig, RFD3InferenceEngine
 from utils import download_weights
+from utils.pipelines import test_rfd3_from_notebook, unconditional_generation
 
 
 download_weights()
 
-
-
-@spaces.GPU(duration=300)  
-def test_rfd3_from_notebook():
-    # Set seed for reproducibility
-    seed_everything(0)
-
-    # Configure RFD3 inference
-    config = RFD3InferenceConfig(
-        specification={
-            'length': 40,  # Generate 80-residue proteins
-        },
-        diffusion_batch_size=2,  # Generate 2 structures per batch
-    )
-
-    # Initialize engine and run generation
-    try: 
-        model = RFD3InferenceEngine(**config)
-        outputs = model.run(
-            inputs=None,      # None for unconditional generation
-            out_dir=None,     # None to return in memory (no file output)
-            n_batches=1,      # Generate 1 batch
-        )
-        return_str = "RDF3 test passed! Generated structures:\n"
-
-        for idx, data in outputs.items():
-            return_str += f"Batch {idx}: {len(data)} structure(s)\n"
-            for i, struct in enumerate(data):
-                return_str += f"Structure {i+1}: {struct.atom_array.array_length()} Atoms\n"
-                #return_str += struct.atom_array
-        return return_str
-    except Exception as e:
-        return f"Error: {str(e)}"
-
     
 # Gradio UI
 with gr.Blocks(title="RFD3 Test") as demo:
@@ -82,38 +49,7 @@ with gr.Blocks(title="RFD3 Test") as demo:
             minimum=10,
             maximum=200
         )
-    # Configure RFD3 inference
-
-
-    # Initialize engine and run generation
-    @spaces.GPU(duration=300) 
-    def unconditional_generation(num_batches, num_designs_per_batch, length):
-
-        config = RFD3InferenceConfig(
-        specification={
-            'length': length,
-        },
-        diffusion_batch_size=num_designs_per_batch,  # Generate 2 structures per batch
-        )
-        try: 
-            model = RFD3InferenceEngine(**config)
-            outputs = model.run(
-                inputs=None,      # None for unconditional generation
-                out_dir=None,     # None to return in memory (no file output)
-                n_batches=num_batches,      # Generate 1 batch
-            )
-            return_str = "RDF3 test passed! Generated structures:\n"
 
-            for idx, data in outputs.items():
-                return_str += f"Batch {idx}: {len(data)} structure(s)\n"
-                for i, struct in enumerate(data):
-                    return_str += f"Structure {i+1}: {struct.atom_array.array_length()} Atoms\n"
-                    #return_str += struct.atom_array
-            return return_str
-        
-        except Exception as e:
-            return f"Error: {str(e)}"
-    
     gen_btn = gr.Button("Run Unconditional Generation")
     gen_output = gr.Textbox(label="Generation Result")
     gen_btn.click(unconditional_generation, inputs=[num_batches, num_designs_per_batch, length], outputs=gen_output)

utils/__init__.py

@@ -1,3 +1,4 @@
 from utils.download_weights import *
+from utils.pipelines import *
 
 __all__ = ["download_weights"]

utils/download_weights.py

@@ -2,7 +2,8 @@ import subprocess
 import os
 from pathlib import Path
 
-MODELS = ["rfd3", "ligandmpnn", "rf3"]
+#MODELS = ["rfd3", "ligandmpnn", "rf3"]
+MODELS = ["rfd3"]
 
 
 # foundry is a package installed automatically upon Space initialization through the Gradio SDK because it is listed in requirements.txt. 

utils/pipelines.py

@@ -0,0 +1,75 @@
+from rfd3.engine import RFD3InferenceConfig, RFD3InferenceEngine
+import gradio as gr
+from lightning.fabric import seed_everything
+import time
+import os
+
+
+@spaces.GPU(duration=300)  
+def test_rfd3_from_notebook():
+    # Set seed for reproducibility
+    seed_everything(0)
+
+    # Configure RFD3 inference
+    config = RFD3InferenceConfig(
+        specification={
+            'length': 40,  # Generate 80-residue proteins
+        },
+        diffusion_batch_size=2,  # Generate 2 structures per batch
+    )
+
+    # Initialize engine and run generation
+    try: 
+        model = RFD3InferenceEngine(**config)
+        outputs = model.run(
+            inputs=None,      # None for unconditional generation
+            out_dir=None,     # None to return in memory (no file output)
+            n_batches=1,      # Generate 1 batch
+        )
+        return_str = "RDF3 test passed! Generated structures:\n"
+
+        for idx, data in outputs.items():
+            return_str += f"Batch {idx}: {len(data)} structure(s)\n"
+            for i, struct in enumerate(data):
+                return_str += f"Structure {i+1}: {struct.atom_array.array_length()} Atoms\n"
+                #return_str += struct.atom_array
+        return return_str
+    except Exception as e:
+        return f"Error: {str(e)}"
+    
+
+
+# Initialize engine and run generation
+@spaces.GPU(duration=300) 
+def unconditional_generation(num_batches, num_designs_per_batch, length):
+
+    config = RFD3InferenceConfig(
+    specification={
+        'length': length,
+    },
+    diffusion_batch_size=num_designs_per_batch,  # Generate 2 structures per batch
+    )
+
+    session_hash = gr.Request.session_hash
+    time_stamp = time.strftime("%Y-%m-%d-%H-%M-%S")
+    directory = f"./outputs/session_{session_hash}_{time_stamp}"
+    os.makedirs(directory, exist_ok=False)
+
+
+    try: 
+        model = RFD3InferenceEngine(**config)
+        outputs = model.run(
+            inputs=None,      # None for unconditional generation
+            out_dir=directory,     # None to return in memory (no file output)
+            n_batches=num_batches,      # Generate 1 batch
+        )
+
+        #for idx, data in outputs.items():
+        #    return_str += f"Batch {idx}: {len(data)} structure(s)\n"
+        #    for i, struct in enumerate(data):
+        #        return_str += f"Structure {i+1}: {struct.atom_array.array_length()} Atoms\n"
+        #        #return_str += struct.atom_array
+        return directory
+    
+    except Exception as e:
+        return f"Error: {str(e)}"