unconditional generation

app.py

@@ -11,35 +11,8 @@ from rfd3.engine import RFD3InferenceConfig, RFD3InferenceEngine
 from utils import download_weights
 
 
-# foundry is a package installed automatically upon Space initialization through the Gradio SDK because it is listed in requirements.txt. 
-# model weights are however not included in the package and must be downloaded separately.
-# the command "foundry install ..." automatically avoids re-downloading models if they are already present in the cache directory.
-#cmd =  f"foundry install rfd3 ligandmpnn rf3"
-#result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
-#if result.returncode == 0:
-#    print("Models installed successfully.")
-#else:   
-#    print(f"Error installing models: {result.stderr}")
-#    print(result.stdout)
-#    print(result.returncode)
-
 download_weights()
 
-# Run once on startup: Install models if missing
-#checkpoint_dir = Path.home() / ".foundry" / "checkpoints"
-#os.environ["FOUNDRY_CHECKPOINT_DIRS"] = str(checkpoint_dir)
-#
-#def install_models():
-#    """Download rfd3, ligandmpnn, rf3 weights once."""
-#    #models = ["rfd3", "ligandmpnn", "rf3"]
-#    models = ["ligandmpnn"] # let's start with only ligand mpnn for testing
-#    for model in models:
-#        if not (checkpoint_dir / model).exists():
-#            print(f"Installing {model}...")
-#            subprocess.check_call(["foundry", "install", model])
-#    print("All models installed.")
-#
-#install_models()  # Executes on app.py load
 
 
 @spaces.GPU(duration=300)  
@@ -77,13 +50,73 @@ def test_rfd3_from_notebook():
     
 # Gradio UI
 with gr.Blocks(title="RFD3 Test") as demo:
-    gr.Markdown("# RFdiffusion3 (RFD3) Model Checker")
+    gr.Markdown("# RFdiffusion3 (RFD3) for Backbone generation")
     gr.Markdown("Models auto-downloaded on launch. Click to test.")
+
     
     test_btn = gr.Button("Run RFD3 Test")
     output = gr.Textbox(label="Test Result")
     
     test_btn.click(test_rfd3_from_notebook, outputs=output)
 
+    gr.Markdown("Unconditional generation of backbones")
+    with gr.Row():
+        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
+        )
+        length = gr.Number(
+            value=40,
+            label="Length of Protein (number of residues)",
+            precision=0,
+            minimum=10,
+            maximum=200
+        )
+    # Configure RFD3 inference
+    config = RFD3InferenceConfig(
+        specification={
+            'length': length,
+
+        },
+        diffusion_batch_size=num_designs_per_batch,  # Generate 2 structures per batch
+    )
+
+    # Initialize engine and run generation
+    def unconditional_generation(num_batches, num_designs_per_batch, length):
+        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)
+
+
+
 if __name__ == "__main__":
     demo.launch()