fix import bug

src/SongFormer/utils/fetch_pretrained.py

@@ -4,6 +4,7 @@ from tqdm import tqdm
 
 
 def download(url, path):
+    """Download file from url to local path with progress bar."""
     if os.path.exists(path):
         print(f"File already exists, skipping download: {path}")
         return
@@ -25,16 +26,31 @@ def download(url, path):
             bar.update(size)
 
 
-# 根据 https://github.com/minzwon/musicfm 下载预训练模型
-download(
-    "https://huggingface.co/minzwon/MusicFM/resolve/main/msd_stats.json",
-    os.path.join("ckpts", "MusicFM", "msd_stats.json"),
-)
-download(
-    "https://huggingface.co/minzwon/MusicFM/resolve/main/pretrained_msd.pt",
-    os.path.join("ckpts", "MusicFM", "pretrained_msd.pt"),
-)
+def download_all(use_mirror: bool = False):
+    """Download all required checkpoints.
+    
+    Args:
+        use_mirror (bool): If True, use hf-mirror.com (for Mainland China).
+    """
+    base_url = "https://hf-mirror.com" if use_mirror else "https://huggingface.co"
 
-# for Mainland China
-# download('https://hf-mirror.com/minzwon/MusicFM/resolve/main/msd_stats.json', os.path.join("ckpts", "MusicFM", "msd_stats.json"))
-# download('https://hf-mirror.com/minzwon/MusicFM/resolve/main/pretrained_msd.pt', os.path.join("ckpts", "MusicFM", "pretrained_msd.pt"))
+    urls = [
+        (f"{base_url}/minzwon/MusicFM/resolve/main/msd_stats.json",
+         os.path.join("ckpts", "MusicFM", "msd_stats.json")),
+        (f"{base_url}/minzwon/MusicFM/resolve/main/pretrained_msd.pt",
+         os.path.join("ckpts", "MusicFM", "pretrained_msd.pt")),
+        (f"{base_url}/ASLP-lab/SongFormer/resolve/main/SongFormer.safetensors",
+         os.path.join("ckpts", "SongFormer.safetensors")),
+        
+        # The content of safetensors is the same as pt, it is recommended to use safetensors
+        # (f"{base_url}/ASLP-lab/SongFormer/resolve/main/SongFormer.pt",
+        #  os.path.join("ckpts", "SongFormer.pt")),
+    ]
+
+    for url, path in urls:
+        download(url, path)
+
+
+if __name__ == "__main__":
+    # By default, use HuggingFace. If you are in Mainland China, change to True
+    download_all(use_mirror=False)