correct test based on colab notebook

app.py

@@ -1,10 +1,14 @@
 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
+
 
 # Download model weights (skips already-downloaded models automatically)
 # In total, ~6GB (3GB for RFD3, 3GB for RF3, <100MB for MPNN); may take a few minutes depending on your connection speed
@@ -58,7 +62,39 @@ def test_rfd3():
             return f"RFD3 test failed: {result.stderr}"
     except Exception as e:
         return f"Error: {str(e)}"
-    
+
+@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}: {len(struct)} 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:
@@ -68,7 +104,7 @@ with gr.Blocks(title="RFD3 Test") as demo:
     test_btn = gr.Button("Run RFD3 Test")
     output = gr.Textbox(label="Test Result")
     
-    test_btn.click(test_rfd3, outputs=output)
+    test_btn.click(test_rfd3_from_notebook, outputs=output)
 
 if __name__ == "__main__":
     demo.launch()