Update download_weights.py

download_weights.py

@@ -1,8 +1,8 @@
-# download_weights.py - Improved with proper OpenAI Point-E model URL
 import os
 import requests
 import torch
 from tqdm import tqdm
+import sys
 
 def download_file(url, destination):
     """
@@ -13,35 +13,57 @@ def download_file(url, destination):
         return
     
     print(f"Downloading {url} to {destination}")
-    response = requests.get(url, stream=True)
-    total_size_in_bytes = int(response.headers.get('content-length', 0))
-    block_size = 1024  # 1 Kibibyte
-    progress_bar = tqdm(total=total_size_in_bytes, unit='iB', unit_scale=True)
-    
-    with open(destination, 'wb') as file:
-        for data in response.iter_content(block_size):
-            progress_bar.update(len(data))
-            file.write(data)
-    
-    progress_bar.close()
-    if total_size_in_bytes != 0 and progress_bar.n != total_size_in_bytes:
-        print("ERROR, something went wrong")
-    else:
-        print(f"Download complete: {destination}")
+    try:
+        response = requests.get(url, stream=True)
+        response.raise_for_status()  # Raise an exception for bad status codes
+        
+        total_size_in_bytes = int(response.headers.get('content-length', 0))
+        block_size = 1024  # 1 Kibibyte
+        progress_bar = tqdm(total=total_size_in_bytes, unit='iB', unit_scale=True)
+        
+        with open(destination, 'wb') as file:
+            for data in response.iter_content(block_size):
+                progress_bar.update(len(data))
+                file.write(data)
+        
+        progress_bar.close()
+        if total_size_in_bytes != 0 and progress_bar.n != total_size_in_bytes:
+            print("ERROR, something went wrong with the download")
+            return False
+        else:
+            print(f"Download complete: {destination}")
+            return True
+    except requests.exceptions.RequestException as e:
+        print(f"Error downloading file: {e}")
+        return False
 
 def main():
     # Create directory for model storage
-    os.makedirs("/tmp/point_e_models", exist_ok=True)
+    model_dir = "/tmp/point_e_models"
+    os.makedirs(model_dir, exist_ok=True)
+    
+    # Ensure directory is writable
+    try:
+        os.chmod(model_dir, 0o777)
+    except Exception as e:
+        print(f"Warning: Could not set permissions on {model_dir}: {e}")
     
-    # URL to download the model weights from Hugging Face Hub
-    # Point-E model weights are available on Hugging Face
+    # URL to download the model weights - use Hugging Face Hub URL
     model_url = "https://huggingface.co/openai/point-e/resolve/main/base40M-textvec.pt"
     
     # Destination path
     model_path = "/tmp/point_e_models/base40M-textvec.pt"
     
-    # Download the model weights
-    download_file(model_url, model_path)
+    # Attempt download 
+    success = download_file(model_url, model_path)
+    if not success:
+        print("Failed to download model weights. Trying alternative source...")
+        # Try an alternative source if the first one fails
+        alt_model_url = "https://github.com/openai/point-e/releases/download/v0.1.0/base40M-textvec.pt"
+        success = download_file(alt_model_url, model_path)
+        if not success:
+            print("Failed to download model weights from alternative source.")
+            sys.exit(1)
     
     # Verify the download
     try:
@@ -49,6 +71,7 @@ def main():
         print(f"Successfully loaded model state dict with {len(state_dict)} keys")
     except Exception as e:
         print(f"Error loading model: {e}")
+        sys.exit(1)
 
 if __name__ == "__main__":
     main()