Improve download_assets.py for Google Drive handling

download_assets.py

@@ -8,7 +8,7 @@ def get_gdrive_file_id(url):
     match = re.search(r"d/([a-zA-Z0-9_-]+)", url) or re.search(r"id=([a-zA-Z0-9_-]+)", url)
     if match:
         return match.group(1)
-    return None
+    return url  # If no match, assume it's already a file ID
 
 def download_gdrive_file(file_id, destination, force=False):
     """Download a file from Google Drive with support for large files"""
@@ -19,52 +19,69 @@ def download_gdrive_file(file_id, destination, force=False):
     # Make the directory if it doesn't exist
     os.makedirs(os.path.dirname(destination), exist_ok=True)
     
-    # Setup URL
+    # First, try the direct download URL
     url = f"https://drive.google.com/uc?export=download&id={file_id}"
     
     # Set up a session to handle cookies
     session = requests.Session()
     
-    try:
-        # First request to get the confirmation token for large files
+    # 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)
-        
-        # 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
     
-    except Exception as e:
-        print(f"Error downloading {destination}: {e}")
-        return False
+    # 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
 
-def download_test_file():
-    """Download a small test file to verify the download functionality"""
-    # Create test directories
+def download_all_assets():
+    """Download all required assets"""
+    # Asset file IDs from Google Drive
+    assets = {
+        "ckpts/landscape.pt": "1-oTNIjCZq3_mGI1XRfzDyCnmjXCvd0Vh",
+        "ckpts/greatest_hits.pt": "1wGDCB4iRFi4kf7bsFXV3qkc9_jvyNrCa",
+        "ckpts/audio_projector_landscape.pth": "1BdjzRJOC8bvyPgrAkJJcCaN3EEJg3STm",
+        "ckpts/audio_projector_gh.pth": "19Uk68PXVOjE3TJl86H-IlMaM1URhU33a",
+        "ckpts/CLAP_weights_2022.pth": "1VK22jxHkFwpxknxQBLd6kIgO5WxQdLFP",
+        "assets/fire_crackling.wav": "1vOAZcbkpo_hre2g26n--lUXdwbTQp22k",
+        "assets/plastic_bag.wav": "15igeDor7a47a-oluSCfO6GeUvFVl2ttb"
+    }
+    
+    # Create directories
     os.makedirs("assets", exist_ok=True)
+    os.makedirs("ckpts", exist_ok=True)
     
-    # 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
+    # Download each asset
+    results = []
+    for destination, file_id in assets.items():
+        if not os.path.exists(destination):
+            try:
+                success = download_gdrive_file(file_id, destination)
+                results.append(f"{destination}: {'Success' if success else 'Failed'}")
+            except Exception as e:
+                results.append(f"{destination}: Error - {str(e)}")
+        else:
+            results.append(f"{destination}: Already exists")
     
-    result = download_gdrive_file(test_file_id, "assets/test_download.wav")
-    return result
+    return results
 
 if __name__ == "__main__":
-    download_test_file()
+    results = download_all_assets()
+    for result in results:
+        print(result)