Update download_models.py

download_models.py

@@ -2,71 +2,77 @@ import os
 from huggingface_hub import hf_hub_download
 from pathlib import Path
 
-def download_with_retry(repo_id, filename, local_dir, max_retries=3):
-    """Download file with retry logic"""
-    for attempt in range(max_retries):
-        try:
-            print(f"Downloading {filename} (attempt {attempt + 1}/{max_retries})...")
-            file_path = hf_hub_download(
-                repo_id=repo_id,
-                filename=filename,
-                local_dir=local_dir,
-                local_dir_use_symlinks=False  # Important for Spaces
-            )
-            print(f"✓ Successfully downloaded: {filename}")
-            return file_path
-        except Exception as e:
-            if attempt == max_retries - 1:
-                print(f"✗ Failed to download {filename}: {e}")
-                raise
-            print(f"Retry {attempt + 1} failed, trying again...")
-    return None
+# Define persistent storage paths
+PERSISTENT_DIR = Path("/data")
+CACHE_DIR = PERSISTENT_DIR / ".cache" / "magicarticulate"
+CACHE_DIR.mkdir(parents=True, exist_ok=True)
 
-def main():
-    print("=" * 50)
-    print("Downloading MagicArticulate Model Checkpoints")
-    print("=" * 50)
+def download_and_link(repo_id, filename, target_local_path):
+    """Download to persistent storage and symlink to app directory"""
+    print(f"Processing {filename}...")
     
-    # Create directories
-    Path("skeleton_ckpt").mkdir(exist_ok=True)
-    Path("third_partys/Michelangelo/checkpoints/aligned_shape_latents").mkdir(
-        parents=True, exist_ok=True
-
-    )
+    # Download to persistent cache
+    # We use hf_hub_download with a specific local_dir in persistent storage
+    # to ensure we control exactly where the file lives.
+    persistent_file_path = CACHE_DIR / filename
     
-    # Download Michelangelo checkpoint (required dependency)
-    print("\n[1/3] Downloading Michelangelo checkpoint...")
-    try:
-        download_with_retry(
-            repo_id="Maikou/Michelangelo",
-            filename="checkpoints/aligned_shape_latents/shapevae-256.ckpt",
-            local_dir="third_partys/Michelangelo"
+    if persistent_file_path.exists():
+        print(f"✓ Found in cache: {persistent_file_path}")
+    else:
+        print(f"Downloading to persistent storage: {persistent_file_path}")
+        hf_hub_download(
+            repo_id=repo_id,
+            filename=filename,
+            local_dir=CACHE_DIR,  # Download directly into /data
+            local_dir_use_symlinks=False
         )
-    except Exception as e:
-        print(f"Warning: Michelangelo download failed: {e}")
-        print("This may affect some features.")
     
-    # Download MagicArticulate spatial checkpoint
-    # print("\n[2/3] Downloading MagicArticulate spatial checkpoint...")
-    # download_with_retry(
-    #     repo_id="Seed3D/MagicArticulate",
-    #     filename="skeleton_ckpt/checkpoint_trainonv2_spatial.pth",
-    #     local_dir="."
-
+    # Create symlink in the application directory
+    target_path = Path(target_local_path)
+    target_path.parent.mkdir(parents=True, exist_ok=True)
+    
+    if target_path.exists():
+        if target_path.is_symlink() and target_path.resolve() == persistent_file_path:
+             print(f"✓ Symlink already exists: {target_path}")
+        else:
+             # If it exists but isn't our symlink, remove it and link
+             if target_path.is_symlink() or target_path.is_file():
+                 target_path.unlink()
+             target_path.symlink_to(persistent_file_path)
+             print(f"✓ Created symlink: {target_path} -> {persistent_file_path}")
+    else:
+        target_path.symlink_to(persistent_file_path)
+        print(f"✓ Created symlink: {target_path} -> {persistent_file_path}")
 
-    # )
+def main():
+    print("=" * 50)
+    print("Setting up models with Persistent Storage")
+    print("=" * 50)
+    
+    # 1. Michelangelo Checkpoint
+    download_and_link(
+        repo_id="Maikou/Michelangelo",
+        filename="checkpoints/aligned_shape_latents/shapevae-256.ckpt",
+        target_local_path="third_partys/Michelangelo/checkpoints/aligned_shape_latents/shapevae-256.ckpt"
+    )
     
-    # Download MagicArticulate hierarchical checkpoint
-    print("\n[3/3] Downloading MagicArticulate hierarchical checkpoint...")
-    download_with_retry(
+    # 2. MagicArticulate Hierarchical Checkpoint
+    download_and_link(
         repo_id="Seed3D/MagicArticulate",
         filename="skeleton_ckpt/checkpoint_trainonv2_hier.pth",
-        local_dir="."
+        target_local_path="skeleton_ckpt/checkpoint_trainonv2_hier.pth"
     )
+
+    # 3. MagicArticulate Spatial Checkpoint (Optional, uncomment if needed)
+    # download_and_link(
+    #     repo_id="Seed3D/MagicArticulate",
+    #     filename="skeleton_ckpt/checkpoint_trainonv2_spatial.pth",
+    #     target_local_path="skeleton_ckpt/checkpoint_trainonv2_spatial.pth"
+    # )
     
     print("\n" + "=" * 50)
-    print("✓ All downloads completed successfully!")
+    print("✓ Setup complete!")
     print("=" * 50)
 
 if __name__ == "__main__":
-    main()
+    main()