modularize weight download

app.py

@@ -8,29 +8,24 @@ import spaces
 from atomworks.io.utils.visualize import view
 from lightning.fabric import seed_everything
 from rfd3.engine import RFD3InferenceConfig, RFD3InferenceEngine
+from utils import download_weights
 
 
-# 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
+# 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()
 
-cmd =  f"foundry install rfd3 ligandmpnn rf3"
-print("Global PATH:", os.environ.get("PATH"))
-
-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_dir = "./models/"
-#if not os.path.exists(download_dir):
-#    cmd  = "foundry install rfd3 ligandmpnn rf3"
-#    result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
-#
-## Run once on startup: Install models if missing
+# Run once on startup: Install models if missing
 #checkpoint_dir = Path.home() / ".foundry" / "checkpoints"
 #os.environ["FOUNDRY_CHECKPOINT_DIRS"] = str(checkpoint_dir)
 #
@@ -46,26 +41,6 @@ else:
 #
 #install_models()  # Executes on app.py load
 
-@spaces.GPU(duration=300)  
-def test_rfd3():
-    """Run a quick rfd3 test design (minimal monomer, 1 step)."""
-    try:
-        cmd = [
-            "rfd3",
-            "design",
-            "--seed", "42",
-            "contigmap.contigs=[A10]",  # Tiny 10-res monomer
-            "--num_designs", "1",
-            "inference.output_prefix=test_output",
-            "--inference.num_diffusion_steps", "10"  # Fast test
-        ]
-        result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
-        if result.returncode == 0:
-            return "RFD3 test passed! Check test_output.pdb. Logs:\n" + result.stdout[-500:]
-        else:
-            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():

utils/__init__.py

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

utils/download_weights.py

@@ -0,0 +1,30 @@
+import subprocess
+import os
+from pathlib import Path
+
+MODELS = ["rfd3", "ligandmpnn", "rf3"]
+
+
+# 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.
+# we however manually check for debugging purposes.
+
+def download_weights():
+    """Download model weights using foundry CLI, skipping already-downloaded models."""
+
+    checkpoint_dir = Path.home() / ".foundry" / "checkpoints"
+    for model in MODELS:
+        model_path = os.path.join(checkpoint_dir, model+".ckpt")
+        if model_path.exists():
+            print(f"{model} already exists at {model_path}, skipping download.")
+        else:
+            cmd = f"foundry install {model} --checkpoint-dir {checkpoint_dir}"
+            print(f"Installing {model}...")
+            result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
+            if result.returncode == 0:
+                print(f"{model} installed successfully.")
+            else:
+                print(f"Error installing {model}: {result.stderr}")
+                print(result.stdout)
+                print(result.returncode)