Fix model downloads

Author: prafullasd

jukebox/make_models.py

@@ -8,7 +8,7 @@
 import torch as t
 import jukebox.utils.dist_adapter as dist
 from jukebox.hparams import Hyperparams, setup_hparams
-from jukebox.utils.gcs_utils import download
+from jukebox.utils.remote_utils import download
 from jukebox.utils.torch_utils import freeze_model
 from jukebox.utils.dist_utils import print_all
 from jukebox.vqvae.vqvae import calculate_strides
@@ -23,15 +23,16 @@
 
 def load_checkpoint(path):
     restore = path
-    if restore[:5] == 'gs://':
-        gs_path = restore
-        local_path = os.path.join(os.path.expanduser("~/.cache"), gs_path[5:])
+    remote_prefix = 'https://openaipublic.blob.core.windows.net/'
+    if restore.startswith(remote_prefix):
+        remote_path = restore
+        local_path = os.path.join(os.path.expanduser("~/.cache"), remote_path[len(remote_prefix):])
         if dist.get_rank() % 8 == 0:
-            print("Downloading from gce")
+            print("Downloading from azure")
             if not os.path.exists(os.path.dirname(local_path)):
                 os.makedirs(os.path.dirname(local_path))
             if not os.path.exists(local_path):
-                download(gs_path, local_path)
+                download(remote_path, local_path)
         restore = local_path
     dist.barrier()
     checkpoint = t.load(restore, map_location=t.device('cpu'))

jukebox/utils/gcs_utils.py renamed to jukebox/utils/remote_utils.py

@@ -1,8 +1,15 @@
-import os
 import sys
 import subprocess
-from time import time
 
+def download(remote_path, local_path, async_download=False):
+    args = ['wget', '-O', local_path, remote_path]
+    print("Running ", " ".join(args))
+    if async_download:
+        subprocess.Popen(args)
+    else:
+        subprocess.call(args)
+
+# GCE
 def gs_download(gs_path, local_path, async_download=False):
     args = ['gsutil',
             '-o', 'GSUtil:parallel_thread_count=1',
@@ -27,16 +34,9 @@ def gs_upload(local_path, gs_path, async_upload=False):
     else:
         subprocess.call(args)
 
-def download(gs_path, local_path, async_download=False):
-    remote_path = gs_path.replace("gs://", "https://storage.googleapis.com/")
-    args = ['wget', '-q', '-O', local_path, remote_path]
-    if async_download:
-        subprocess.Popen(args)
-    else:
-        subprocess.call(args)
-
 def ls(regex):
     outputs = subprocess.check_output(['gsutil', 'ls', regex]).decode(sys.stdout.encoding)
     outputs = outputs.split('\n')
     outputs = [output for output in outputs if output is not '']
     return outputs
+