committing...

app.py

@@ -14,7 +14,7 @@ import os
 # Global variables for model and tokenizer
 global_generator = None
 global_tokenizer = None
-default_audio_path = "testingtesting.wav"  # Your default audio file
+default_audio_path = "sample.wav"  # Changed from "testingtesting.wav"
 
 def init_model(use_pure_audio_ablation: bool = False) -> Tuple[nn.Module, object]:
     """Initialize the model and tokenizer"""

utils/dist.py

@@ -8,6 +8,7 @@ import requests
 import hashlib
 
 from io import BytesIO  
+from huggingface_hub import hf_hub_download
 
 def rank0():
     rank = os.environ.get('RANK')
@@ -75,17 +76,12 @@ def init_dist():
     return rank, local_rank, world_size
 
 def load_ckpt(load_from_location, expected_hash=None):
+    os.environ['HF_HUB_ENABLE_HF_TRANSFER'] = '1' #Disable this to speed up debugging errors with downloading from the hub
     if local0():
-        os.makedirs('ckpt', exist_ok=True)
-        url = f"https://ckpt.si.inc/hertz-dev/{load_from_location}.pt"
-        save_path = f"ckpt/{load_from_location}.pt"
-        if not os.path.exists(save_path):
-            response = requests.get(url, stream=True)
-            total_size = int(response.headers.get('content-length', 0))
-            with open(save_path, 'wb') as f, tqdm(total=total_size, desc=f'Downloading {load_from_location}.pt', unit='GB', unit_scale=1/(1024*1024*1024)) as pbar:
-                for chunk in response.iter_content(chunk_size=8192):
-                    f.write(chunk)
-                    pbar.update(len(chunk))
+        repo_id = "si-pbc/hertz-dev"
+        print0(f'Loading checkpoint from repo_id {repo_id} and filename {load_from_location}.pt. This may take a while...')
+        save_path = hf_hub_download(repo_id=repo_id, filename=f"{load_from_location}.pt")
+        print0(f'Downloaded checkpoint to {save_path}')
         if expected_hash is not None:
             with open(save_path, 'rb') as f:
                 file_hash = hashlib.md5(f.read()).hexdigest()
@@ -94,6 +90,9 @@ def load_ckpt(load_from_location, expected_hash=None):
                 os.remove(save_path)
                 return load_ckpt(load_from_location, expected_hash)
     if T.distributed.is_initialized():
-        T.distributed.barrier() # so that ranks don't try to load checkpoint before it's finished downloading
-    loaded = T.load(f"ckpt/{load_from_location}.pt", weights_only=False, map_location='cpu')    
+        save_path = [save_path]
+        T.distributed.broadcast_object_list(save_path, src=0)
+        save_path = save_path[0]
+    loaded = T.load(save_path, weights_only=False, map_location='cpu')
+    print0(f'Loaded checkpoint from {save_path}')
     return loaded