Add simplified download_assets.py

download_assets.py

@@ -1,7 +1,7 @@
 import os
-import re
 import requests
 from tqdm import tqdm
+import re
 
 def get_gdrive_file_id(url):
     """Extract file ID from Google Drive URL"""
@@ -10,75 +10,61 @@ def get_gdrive_file_id(url):
         return match.group(1)
     return None
 
-def download_gdrive_file(file_id, destination):
+def download_gdrive_file(file_id, destination, force=False):
     """Download a file from Google Drive with support for large files"""
-    
-    if os.path.exists(destination):
+    if os.path.exists(destination) and not force:
         print(f"File already exists: {destination}")
-        return
+        return True
     
     # Make the directory if it doesn't exist
     os.makedirs(os.path.dirname(destination), exist_ok=True)
     
-    # First, try the direct download URL
+    # Setup URL
     url = f"https://drive.google.com/uc?export=download&id={file_id}"
     
     # Set up a session to handle cookies
     session = requests.Session()
     
-    # First request to get the confirmation token for large files
-    response = session.get(url, stream=True)
-    
-    # Check if there's a download confirmation page
-    if "confirm" in response.url:
-        # Extract confirmation token
-        token = response.url.split("confirm=")[1].split("&")[0]
-        url = f"{url}&confirm={token}"
+    try:
+        # First request to get the confirmation token for large files
         response = session.get(url, stream=True)
+        
+        # Check if there's a download confirmation page
+        if "confirm" in response.url:
+            # Extract confirmation token
+            token = response.url.split("confirm=")[1].split("&")[0]
+            url = f"{url}&confirm={token}"
+            response = session.get(url, stream=True)
+        
+        # Get file size for progress bar
+        total_size = int(response.headers.get('content-length', 0))
+        
+        # Download the file with progress bar
+        print(f"Downloading to {destination} ({total_size/(1024*1024):.1f} MB)...")
+        with open(destination, 'wb') as f:
+            with tqdm(total=total_size, unit='B', unit_scale=True) as pbar:
+                for chunk in response.iter_content(chunk_size=1024*1024):
+                    if chunk:
+                        f.write(chunk)
+                        pbar.update(len(chunk))
+        
+        print(f"Downloaded {destination} successfully!")
+        return True
     
-    # Get file size for progress bar
-    total_size = int(response.headers.get('content-length', 0))
-    
-    # Download the file with progress bar
-    print(f"Downloading to {destination} ({total_size/(1024*1024):.1f} MB)...")
-    with open(destination, 'wb') as f:
-        with tqdm(total=total_size, unit='B', unit_scale=True) as pbar:
-            for chunk in response.iter_content(chunk_size=1024*1024):
-                if chunk:
-                    f.write(chunk)
-                    pbar.update(len(chunk))
-    
-    print(f"Downloaded {destination} successfully!")
+    except Exception as e:
+        print(f"Error downloading {destination}: {e}")
+        return False
 
-def main():
-    """Download all required assets"""
-    # Original sharing URLs
-    gdrive_urls = {
-        "assets/fire_crackling.wav": "https://drive.google.com/file/d/1vOAZcbkpo_hre2g26n--lUXdwbTQp22k/view?usp=drive_link",
-        "assets/plastic_bag.wav": "https://drive.google.com/file/d/15igeDor7a47a-oluSCfO6GeUvFVl2ttb/view?usp=sharing",
-        "ckpts/landscape.pt": "https://drive.google.com/file/d/1-oTNIjCZq3_mGI1XRfzDyCnmjXCvd0Vh/view?usp=drive_link",
-        "ckpts/greatest_hits.pt": "https://drive.google.com/file/d/1wGDCB4iRFi4kf7bsFXV3qkc9_jvyNrCa/view?usp=drive_link",
-        "ckpts/audio_projector_landscape.pth": "https://drive.google.com/file/d/1BdjzRJOC8bvyPgrAkJJcCaN3EEJg3STm/view?usp=sharing",
-        "ckpts/audio_projector_gh.pth": "https://drive.google.com/file/d/19Uk68PXVOjE3TJl86H-IlMaM1URhU33a/view?usp=sharing",
-        "ckpts/CLAP_weights_2022.pth": "https://drive.google.com/file/d/1VK22jxHkFwpxknxQBLd6kIgO5WxQdLFP/view?usp=sharing"
-    }
-    
-    # Create necessary directories
+def download_test_file():
+    """Download a small test file to verify the download functionality"""
+    # Create test directories
     os.makedirs("assets", exist_ok=True)
-    os.makedirs("ckpts", exist_ok=True)
     
-    # Download each file
-    for destination, url in gdrive_urls.items():
-        file_id = get_gdrive_file_id(url)
-        if file_id:
-            try:
-                download_gdrive_file(file_id, destination)
-            except Exception as e:
-                print(f"Error downloading {destination}: {e}")
-        else:
-            print(f"Could not extract file ID from {url}")
+    # Small test file - this is a public file ID from a small test image
+    test_file_id = "1vOAZcbkpo_hre2g26n--lUXdwbTQp22k" # One of your fire_crackling.wav ID
     
-    print("Download process completed!")
+    result = download_gdrive_file(test_file_id, "assets/test_download.wav")
+    return result
 
 if __name__ == "__main__":
-    main()
+    download_test_file()