download ckp updated

download_ckp.py

@@ -3,13 +3,23 @@ import os
 from pathlib import Path
 
 ckp = Path("/app/ckp")
-if not list(ckp.glob("*.pth")):
+ckp_single_cell = ckp / "single_cell"
+ckp_spheroid = ckp / "spheroid"
+has_any = list(ckp.glob("*.pth")) or list(ckp_single_cell.glob("*.pth")) or list(ckp_spheroid.glob("*.pth"))
+
+if not has_any:
     try:
         from huggingface_hub import hf_hub_download, list_repo_files
 
         repo = os.environ.get("HF_MODEL_REPO", "kaveh/Shape2Force")
         files = list_repo_files(repo)
         pth_files = [f for f in files if f.startswith("ckp/") and f.endswith(".pth")]
+        # For spheroid: only download ckp_spheroid_FN.pth (not ckp_spheroid_GN.pth or others)
+        def should_download(f):
+            if "spheroid" in f and "ckp_spheroid_FN.pth" not in f:
+                return False
+            return True
+        pth_files = [f for f in pth_files if should_download(f)]
         for f in pth_files:
             hf_hub_download(repo_id=repo, filename=f, local_dir="/app")
         print("Downloaded checkpoints from", repo)