Update to CDN for faster download

openai · Nov 14, 2020 · 08efbbc · adassesans · Jul 15, 2023 · 08efbbc
1 parent 1953375
commit 08efbbc
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 10 deletions.
diff --git a/jukebox/hparams.py b/jukebox/hparams.py
@@ -33,6 +33,8 @@ def setup_hparams(hparam_set_names, kwargs):
 )
 HPARAMS_REGISTRY["easy"] = easy
 
+REMOTE_PREFIX = 'https://openaipublic.azureedge.net/'
+
 # Model hps
 vqvae = Hyperparams(
     levels = 3,
@@ -52,7 +54,7 @@ def setup_hparams(hparam_set_names, kwargs):
     depth = 4,
     m_conv = 1.0,
     dilation_growth_rate = 3,
-    restore_vqvae=f'https://openaipublic.blob.core.windows.net/jukebox/models/5b/vqvae.pth.tar',
+    restore_vqvae=REMOTE_PREFIX + 'jukebox/models/5b/vqvae.pth.tar',
 )
 HPARAMS_REGISTRY["vqvae"] = vqvae
 
@@ -85,15 +87,15 @@ def setup_hparams(hparam_set_names, kwargs):
 
 upsampler_level_0 = Hyperparams(
     level=0,
-    restore_prior='https://openaipublic.blob.core.windows.net/jukebox/models/5b/prior_level_0.pth.tar'
+    restore_prior=REMOTE_PREFIX + 'jukebox/models/5b/prior_level_0.pth.tar'
 )
 upsampler_level_0.update(upsamplers)
 HPARAMS_REGISTRY["upsampler_level_0"] = upsampler_level_0
 
 upsampler_level_1 = Hyperparams(
     level=1,
     cond_res_scale=True,
-    restore_prior='https://openaipublic.blob.core.windows.net/jukebox/models/5b/prior_level_1.pth.tar'
+    restore_prior=REMOTE_PREFIX + 'jukebox/models/5b/prior_level_1.pth.tar'
 )
 upsampler_level_1.update(upsamplers)
 HPARAMS_REGISTRY["upsampler_level_1"] = upsampler_level_1
@@ -115,7 +117,7 @@ def setup_hparams(hparam_set_names, kwargs):
     n_tokens=0,
     prime_loss_fraction=0.0,
     merged_decoder=True,
-    restore_prior='https://openaipublic.blob.core.windows.net/jukebox/models/5b/prior_level_2.pth.tar',
+    restore_prior=REMOTE_PREFIX + 'jukebox/models/5b/prior_level_2.pth.tar',
     fp16_params=True,
 )
 prior_5b.update(labels)
@@ -145,7 +147,7 @@ def setup_hparams(hparam_set_names, kwargs):
     n_tokens=512,
     prime_loss_fraction=0.4,
     merged_decoder=True,
-    restore_prior='https://openaipublic.blob.core.windows.net/jukebox/models/5b_lyrics/prior_level_2.pth.tar',
+    restore_prior=REMOTE_PREFIX + 'jukebox/models/5b_lyrics/prior_level_2.pth.tar',
     fp16_params=True,
     alignment_layer=68,
     alignment_head=2,
@@ -177,7 +179,7 @@ def setup_hparams(hparam_set_names, kwargs):
     n_tokens=384,
     prime_loss_fraction=0.4,
     single_enc_dec=True,
-    restore_prior='https://openaipublic.blob.core.windows.net/jukebox/models/1b_lyrics/prior_level_2.pth.tar',
+    restore_prior=REMOTE_PREFIX + 'jukebox/models/1b_lyrics/prior_level_2.pth.tar',
     fp16_params=False,
     alignment_layer=63,
     alignment_head=0,

diff --git a/jukebox/make_models.py b/jukebox/make_models.py
@@ -7,7 +7,7 @@
 import numpy as np
 import torch as t
 import jukebox.utils.dist_adapter as dist
-from jukebox.hparams import Hyperparams, setup_hparams
+from jukebox.hparams import Hyperparams, setup_hparams, REMOTE_PREFIX
 from jukebox.utils.remote_utils import download
 from jukebox.utils.torch_utils import freeze_model
 from jukebox.utils.dist_utils import print_all
@@ -23,10 +23,9 @@
 
 def load_checkpoint(path):
     restore = path
-    remote_prefix = 'https://openaipublic.blob.core.windows.net/'
-    if restore.startswith(remote_prefix):
+    if restore.startswith(REMOTE_PREFIX):
         remote_path = restore
-        local_path = os.path.join(os.path.expanduser("~/.cache"), remote_path[len(remote_prefix):])
+        local_path = os.path.join(os.path.expanduser("~/.cache"), remote_path[len(REMOTE_PREFIX):])
         if dist.get_rank() % 8 == 0:
             print("Downloading from azure")
             if not os.path.exists(os.path.dirname(local_path)):