Add gdown support for drive downloads (more reliable)

dorarad · Feb 14, 2022 · 995449a · 995449a
1 parent 7267aab
commit 995449a
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/pytorch_version/dnnlib/util.py b/pytorch_version/dnnlib/util.py
@@ -20,6 +20,7 @@
 import urllib
 import urllib.request
 import uuid
+import gdown
 
 from distutils.util import strtobool
 from typing import Any, List, Tuple, Union
@@ -394,6 +395,10 @@ def open_url(url: str, cache_dir: str = None, num_attempts: int = 10, verbose: b
             filename = cache_files[0]
             return filename if return_filename else open(filename, "rb")
 
+    if "drive.google" in url:
+        filename = gdown.download(url, quiet = False)
+        return open(filename, "rb")
+
     # Download
     url_name = None
     url_data = None