Update download_models.py

download_models.py

@@ -1,36 +1,72 @@
-from huggingface_hub import hf_hub_download
 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
 
-def download_model(repo_id, filename):
-    """Download and cache model using HF Hub"""
-    print(f"Downloading {filename}...")
+def main():
+    print("=" * 50)
+    print("Downloading MagicArticulate Model Checkpoints")
+    print("=" * 50)
     
-    # HF Hub automatically caches in ~/.cache/huggingface
-    file_path = hf_hub_download(
-        repo_id=repo_id,
-        filename=filename,
-        cache_dir=os.environ.get("HF_HOME", "/data/.cache/huggingface"),  # Space persistent storage
+    # Create directories
+    Path("skeleton_ckpt").mkdir(exist_ok=True)
+    Path("third_partys/Michelangelo/checkpoints/aligned_shape_latents").mkdir(
+        parents=True, exist_ok=True
+
     )
     
-    # Create symlink in expected location
-    local_path = filename
-    os.makedirs(os.path.dirname(local_path), exist_ok=True)
+    # 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"
+        )
+    except Exception as e:
+        print(f"Warning: Michelangelo download failed: {e}")
+        print("This may affect some features.")
     
-    if not os.path.exists(local_path):
-        os.symlink(file_path, local_path)
-    
-    print(f"✓ {filename} ready at {local_path}")
-    return file_path
+    # 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="."
+
 
-def main():
-    # # Download spatial checkpoint
-    # download_model(
-    #     "Seed3D/MagicArticulate",
-    #     "skeleton_ckpt/checkpoint_trainonv2_spatial.pth"
     # )
     
-    # Download hierarchical checkpoint  
-    download_model(
-        "Seed3D/MagicArticulate",
-        "skeleton_ckpt/checkpoint_trainonv2_hier.pth"
+    # Download MagicArticulate hierarchical checkpoint
+    print("\n[3/3] Downloading MagicArticulate hierarchical checkpoint...")
+    download_with_retry(
+        repo_id="Seed3D/MagicArticulate",
+        filename="skeleton_ckpt/checkpoint_trainonv2_hier.pth",
+        local_dir="."
     )
+    
+    print("\n" + "=" * 50)
+    print("✓ All downloads completed successfully!")
+    print("=" * 50)
+
+if __name__ == "__main__":
+    main()