update bundle download api

Signed-off-by: Yiheng Wang <[email protected]>

Author: yiheng-wang-nv

monai/bundle/scripts.py

@@ -600,11 +600,15 @@ def download(
             _download_from_github(repo=repo_, download_path=bundle_dir_, filename=name_ver, progress=progress_)
         elif source_ == "monaihosting":
             try:
+                extract_path = os.path.join(bundle_dir_, name_)
+                huggingface_hub.snapshot_download(repo_id=f"MONAI/{name_}", revision=version_, local_dir=extract_path)
+            except (huggingface_hub.errors.RevisionNotFoundError, huggingface_hub.errors.RepositoryNotFoundError):
+                # if bundle or version not found from huggingface, download from ngc monaihosting
                 _download_from_monaihosting(
                     download_path=bundle_dir_, filename=name_, version=version_, progress=progress_
                 )
             except urllib.error.HTTPError:
-                # for monaihosting bundles, if cannot download from default host, download according to bundle_info
+                # if also cannot download from ngc monaihosting, download according to bundle_info
                 _download_from_bundle_info(
                     download_path=bundle_dir_, filename=name_, version=version_, progress=progress_
                 )

monai/utils/jupyter_utils.py

@@ -234,7 +234,7 @@ def plot_engine_status(
 
 
 def _get_loss_from_output(
-    output: list[torch.Tensor | dict[str, torch.Tensor]] | dict[str, torch.Tensor] | torch.Tensor,
+    output: list[torch.Tensor | dict[str, torch.Tensor]] | dict[str, torch.Tensor] | torch.Tensor
 ) -> torch.Tensor:
     """Returns a single value from the network output, which is a dict or tensor."""
 

tests/bundle/test_bundle_download.py

@@ -63,6 +63,8 @@
 
 TEST_CASE_6 = [["models/model.pt", "configs/train.json"], "renalStructures_CECT_segmentation", "0.1.0"]
 
+TEST_CASE_6_HF = [["models/model.pt", "configs/train.yaml"], "mednist_ddpm", "1.0.1"]
+
 TEST_CASE_7 = [
     ["model.pt", "model.ts", "network.json", "test_output.pt", "test_input.pt"],
     "test_bundle",
@@ -193,6 +195,7 @@ def test_ngc_private_source_download_bundle(self, bundle_files, bundle_name, _ur
 
     @parameterized.expand([TEST_CASE_6])
     @skip_if_quick
+    @skipUnless(has_huggingface_hub, "Requires `huggingface_hub`.")
     def test_monaihosting_source_download_bundle(self, bundle_files, bundle_name, version):
         with skip_if_downloading_fails():
             # download a single file from url, also use `args_file`
@@ -239,6 +242,7 @@ def test_list_latest_versions(self):
         self.assertEqual(_list_latest_versions(data), ["1.1", "1.0"])
 
     @skip_if_quick
+    @skipUnless(has_huggingface_hub, "Requires `huggingface_hub`.")
     @patch("monai.bundle.scripts.get_versions", return_value={"version": "1.2"})
     def test_download_monaihosting(self, mock_get_versions):
         """Test checking MONAI version from a metadata file."""
@@ -333,6 +337,7 @@ def test_load_weights(self, bundle_files, bundle_name, repo, device, model_file)
 
     @parameterized.expand([TEST_CASE_8])
     @skip_if_quick
+    @skipUnless(has_huggingface_hub, "Requires `huggingface_hub`.")
     def test_load_weights_with_net_override(self, bundle_name, device, net_override):
         with skip_if_downloading_fails():
             # download bundle, and load weights from the downloaded path