4042 Hugging Face Hub integration

docs/requirements.txt

@@ -37,3 +37,4 @@ optuna
 opencv-python-headless
 onnx>=1.13.0
 onnxruntime; python_version <= '3.10'
+huggingface_hub

docs/source/installation.md

@@ -10,8 +10,8 @@
 		- [Uninstall the packages](#uninstall-the-packages)
 	- [From conda-forge](#from-conda-forge)
 	- [From GitHub](#from-github)
-		- [Option 1 (as a part of your system-wide module)](#option-1-as-a-part-of-your-system-wide-module)
-		- [Option 2 (editable installation)](#option-2-editable-installation)
+		- [Option 1 (as a part of your system-wide module):](#option-1-as-a-part-of-your-system-wide-module)
+		- [Option 2 (editable installation):](#option-2-editable-installation)
 	- [Validating the install](#validating-the-install)
 	- [MONAI version string](#monai-version-string)
 	- [From DockerHub](#from-dockerhub)
@@ -254,10 +254,10 @@ Since MONAI v0.2.0, the extras syntax such as `pip install 'monai[nibabel]'` is
 - The options are
 
 ```
-[nibabel, skimage, pillow, tensorboard, gdown, ignite, torchvision, itk, tqdm, lmdb, psutil, cucim, openslide, pandas, einops, transformers, mlflow, clearml, matplotlib, tensorboardX, tifffile, imagecodecs, pyyaml, fire, jsonschema, ninja, pynrrd, pydicom, h5py, nni, optuna, onnx, onnxruntime]
+[nibabel, skimage, pillow, tensorboard, gdown, ignite, torchvision, itk, tqdm, lmdb, psutil, cucim, openslide, pandas, einops, transformers, mlflow, clearml, matplotlib, tensorboardX, tifffile, imagecodecs, pyyaml, fire, jsonschema, ninja, pynrrd, pydicom, h5py, nni, optuna, onnx, onnxruntime, huggingface_hub]
 ```
 
 which correspond to `nibabel`, `scikit-image`, `pillow`, `tensorboard`,
-`gdown`, `pytorch-ignite`, `torchvision`, `itk`, `tqdm`, `lmdb`, `psutil`, `cucim`, `openslide-python`, `pandas`, `einops`, `transformers`, `mlflow`, `clearml`, `matplotlib`, `tensorboardX`, `tifffile`, `imagecodecs`, `pyyaml`, `fire`, `jsonschema`, `ninja`, `pynrrd`, `pydicom`, `h5py`, `nni`, `optuna`, `onnx`, `onnxruntime`, respectively.
+`gdown`, `pytorch-ignite`, `torchvision`, `itk`, `tqdm`, `lmdb`, `psutil`, `cucim`, `openslide-python`, `pandas`, `einops`, `transformers`, `mlflow`, `clearml`, `matplotlib`, `tensorboardX`, `tifffile`, `imagecodecs`, `pyyaml`, `fire`, `jsonschema`, `ninja`, `pynrrd`, `pydicom`, `h5py`, `nni`, `optuna`, `onnx`, `onnxruntime`, `huggingface_hub`, respectively.
 
 - `pip install 'monai[all]'` installs all the optional dependencies.

monai/bundle/scripts.py

@@ -57,6 +57,7 @@
 Checkpoint, has_ignite = optional_import("ignite.handlers", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Checkpoint")
 requests_get, has_requests = optional_import("requests", name="get")
 onnx, _ = optional_import("onnx")
+huggingface_hub, _ = optional_import("huggingface_hub")
 
 logger = get_logger(module_name=__name__)
 
@@ -192,6 +193,13 @@ def _download_from_ngc(
     extractall(filepath=filepath, output_dir=extract_path, has_base=True)
 
 
+def _download_from_huggingface_hub(repo: str, download_path: str, filename: str, version: str) -> None:
+    if len(repo.split("/")) != 2:
+        raise ValueError("if source is `huggingface_hub`, repo should be in the form `repo_owner/repo_name`")
+    extract_path = os.path.join(download_path, filename)
+    huggingface_hub.snapshot_download(repo_id=repo, revision=version, local_dir=extract_path, local_dir_use_symlinks="auto")
+
+
 def _get_latest_bundle_version(source: str, name: str, repo: str) -> dict[str, list[str] | str] | Any | None:
     if source == "ngc":
         name = _add_ngc_prefix(name)
@@ -203,6 +211,10 @@ def _get_latest_bundle_version(source: str, name: str, repo: str) -> dict[str, l
     elif source == "github":
         repo_owner, repo_name, tag_name = repo.split("/")
         return get_bundle_versions(name, repo=f"{repo_owner}/{repo_name}", tag=tag_name)["latest_version"]
+    elif source == "huggingface_hub":
+        huggingface_hub.list_repo_refs(repo_id=f"{repo}/{name}", repo_type="model")
+        #TODO: implement this
+        return None
     else:
         raise ValueError(f"To get the latest bundle version, source should be 'github' or 'ngc', got {source}.")
 
@@ -247,6 +259,9 @@ def download(
         # Execute this module as a CLI entry, and download bundle from ngc with latest version:
         python -m monai.bundle download --name <bundle_name> --source "ngc" --bundle_dir "./"
 
+        # Execute this module as a CLI entry, and download bundle from Hugging Face Hub:
+        python -m monai.bundle download --name "bundle_name" --source "huggingface_hub" --repo "repo_owner/repo_name"
+
         # Execute this module as a CLI entry, and download bundle via URL:
         python -m monai.bundle download --name <bundle_name> --url <url>
 
@@ -265,14 +280,15 @@ def download(
             "monai_brats_mri_segmentation" in ngc:
             https://catalog.ngc.nvidia.com/models?filters=&orderBy=scoreDESC&query=monai.
         version: version name of the target bundle to download, like: "0.1.0". If `None`, will download
-            the latest version.
+            the latest version. If `source` is "huggingface_hub", this argument is a Git revision id.
         bundle_dir: target directory to store the downloaded data.
             Default is `bundle` subfolder under `torch.hub.get_dir()`.
         source: storage location name. This argument is used when `url` is `None`.
             In default, the value is achieved from the environment variable BUNDLE_DOWNLOAD_SRC, and
-            it should be "ngc" or "github".
-        repo: repo name. This argument is used when `url` is `None` and `source` is "github".
-            If used, it should be in the form of "repo_owner/repo_name/release_tag".
+            it should be "ngc", "github", or "huggingface_hub".
+        repo: repo name. This argument is used when `url` is `None` and `source` is "github" or "huggingface_hub".
+            If `source` is "github", it should be in the form of "repo_owner/repo_name/release_tag".
+            If `source` is "huggingface_hub", it should be in the form of "repo_owner/repo_name".
         url: url to download the data. If not `None`, data will be downloaded directly
             and `source` will not be checked.
             If `name` is `None`, filename is determined by `monai.apps.utils._basename(url)`.
@@ -332,9 +348,18 @@ def download(
                 remove_prefix=remove_prefix_,
                 progress=progress_,
             )
+        elif source_ == "huggingface_hub":
+            if name_ is None:
+                raise ValueError(f"To download from source: 'huggingface_hub', `name` must be provided, got {name_}.")
+            _download_from_huggingface_hub(
+                repo=repo_,
+                download_path=bundle_dir_,
+                filename=name_,
+                version=version_,
+            )
         else:
             raise NotImplementedError(
-                f"Currently only download from `url`, source 'github' or 'ngc' are implemented, got source: {source_}."
+                f"Currently only download from `url`, source 'github', 'ngc', or 'huggingface_hub' are implemented, got source: {source_}."
             )
 
 
@@ -365,7 +390,7 @@ def load(
             "monai_brats_mri_segmentation" in ngc:
             https://catalog.ngc.nvidia.com/models?filters=&orderBy=scoreDESC&query=monai.
         version: version name of the target bundle to download, like: "0.1.0". If `None`, will download
-            the latest version.
+            the latest version. If `source` is "huggingface_hub", this argument is a Git revision id.
         model_file: the relative path of the model weights or TorchScript module within bundle.
             If `None`, "models/model.pt" or "models/model.ts" will be used.
         load_ts_module: a flag to specify if loading the TorchScript module.
@@ -374,9 +399,10 @@ def load(
         source: storage location name. This argument is used when `model_file` is not existing locally and need to be
             downloaded first.
             In default, the value is achieved from the environment variable BUNDLE_DOWNLOAD_SRC, and
-            it should be "ngc" or "github".
-        repo: repo name. This argument is used when `url` is `None` and `source` is "github".
-            If used, it should be in the form of "repo_owner/repo_name/release_tag".
+            it should be "ngc", "github", or "huggingface_hub".
+        repo: repo name. This argument is used when `url` is `None` and `source` is "github" or "huggingface_hub".
+            If `source` is "github", it should be in the form of "repo_owner/repo_name/release_tag".
+            If `source` is "huggingface_hub", it should be in the form of "repo_owner/repo_name".
         remove_prefix: This argument is used when `source` is "ngc". Currently, all ngc bundles
             have the ``monai_`` prefix, which is not existing in their model zoo contrasts. In order to
             maintain the consistency between these two sources, remove prefix is necessary.
@@ -1501,3 +1527,16 @@ def init_bundle(
         copyfile(str(ckpt_file), str(models_dir / "model.pt"))
     elif network is not None:
         save_state(network, str(models_dir / "model.pt"))
+
+
+def push_to_hf_hub(bundle_dir: str, repo_name: str) -> None:
+    """
+    Push the current bundle to the Hugging Face Hub.
+
+    Args:
+        bundle_dir: path to the bundle directory to push
+        repo_name: name of the repo to create or push to the HF Hub
+    """
+    hf_api = huggingface_hub.HfApi()
+    repo_id = hf_api.create_repo(name=repo_name, exist_ok=True)
+    return hf_api.upload_folder(path=bundle_dir, repo_id=repo_id)

requirements-dev.txt

@@ -52,3 +52,4 @@ onnx>=1.13.0
 onnxruntime; python_version <= '3.10'
 typeguard<3  # https://github.com/microsoft/nni/issues/5457
 filelock!=3.12.0  # https://github.com/microsoft/nni/issues/5523
+huggingface_hub

setup.cfg

@@ -79,6 +79,7 @@ all =
     optuna
     onnx>=1.13.0
     onnxruntime; python_version <= '3.10'
+    huggingface_hub
 nibabel =
     nibabel
 ninja =
@@ -145,6 +146,8 @@ onnx =
 # # workaround https://github.com/Project-MONAI/MONAI/issues/5882
 # MetricsReloaded =
 #     MetricsReloaded @ git+https://github.com/Project-MONAI/MetricsReloaded@monai-support#egg=MetricsReloaded
+huggingface_hub =
+    huggingface_hub
 
 [flake8]
 select = B,C,E,F,N,P,T4,W,B9