[codex] Trigger Hugging Face access during inference

NV-Segment-CT/configs/batch_inference.json

@@ -1,14 +1,16 @@
 {
     "input_dir": "@bundle_root",
     "input_suffix": "*.nii.gz",
-    "input_root_abs": "$os.path.abspath(@input_dir)",
+    "input_root_abs": "$os.path.realpath(@input_dir)",
+    "output_dir_abs": "$os.path.realpath(@output_dir)",
     "batch_skip_dir_names": [],
     "batch_skip_dir_prefixes": [],
     "batch_resume_skip_existing": true,
     "batch_use_input_list_cache": true,
     "batch_cache_wait_sec": 120,
-    "input_list": "$scripts.batch_inference_utils.build_input_list(os.path.abspath(@input_dir), os.path.abspath(@output_dir), @output_postfix, @output_ext, @batch_skip_dir_names, @batch_skip_dir_prefixes, @batch_resume_skip_existing, @batch_use_input_list_cache, @batch_cache_wait_sec)",
+    "input_list": "$scripts.batch_inference_utils.build_input_list(@input_root_abs, @output_dir_abs, @output_postfix, @output_ext, @batch_skip_dir_names, @batch_skip_dir_prefixes, @batch_resume_skip_existing, @batch_use_input_list_cache, @batch_cache_wait_sec)",
     "input_dicts": "$[{'image': x, 'label_prompt': @everything_labels} for x in @input_list]",
     "dataset#data": "@input_dicts",
+    "postprocessing#transforms#4#output_dir": "@output_dir_abs",
     "postprocessing#transforms#4#data_root_dir": "@input_root_abs"
 }

NV-Segment-CT/configs/inference.json

@@ -9,6 +9,10 @@
         "$import pathlib"
     ],
     "bundle_root": "./",
+    "huggingface_repo_id": "nvidia/NV-Segment-CT",
+    "huggingface_checkpoint_file": "vista3d_pretrained_model/model.pt",
+    "huggingface_download_counter_file": "config.json",
+    "checkpoint_path": "$@bundle_root + '/models/model.pt'",
     "image_key": "image",
     "output_dir": "$@bundle_root + '/eval'",
     "output_ext": ".nii.gz",
@@ -174,7 +178,7 @@
     ],
     "checkpointloader": {
         "_target_": "CheckpointLoader",
-        "load_path": "$@bundle_root + '/models/model.pt'",
+        "load_path": "@checkpoint_path",
         "load_dict": {
             "model": "@network"
         },
@@ -195,6 +199,7 @@
         }
     },
     "initialize": [
+        "$scripts.prepare_huggingface_checkpoint(@huggingface_repo_id, @huggingface_checkpoint_file, @checkpoint_path, @huggingface_download_counter_file)",
         "$monai.utils.set_determinism(seed=123)",
         "$@checkpointloader(@evaluator)"
     ],

NV-Segment-CT/configs/mgpu_inference.json

@@ -18,6 +18,8 @@
         "$import torch.distributed as dist",
         "$dist.is_initialized() or dist.init_process_group(backend='nccl')",
         "$torch.cuda.set_device(@device)",
+        "$scripts.prepare_huggingface_checkpoint(@huggingface_repo_id, @huggingface_checkpoint_file, @checkpoint_path, @huggingface_download_counter_file)",
+        "$dist.barrier()",
         "$monai.utils.set_determinism(seed=123)",
         "$@checkpointloader(@evaluator)"
     ],

NV-Segment-CT/docs/README.md

@@ -14,13 +14,10 @@ conda activate vista3d-nv
 git clone https://github.com/NVIDIA-Medtech/NV-Segment-CTMR.git
 cd NV-Segment-CTMR/NV-Segment-CT;
 pip install -r requirements.txt;
-
-mkdir -p models
-# Option 1: Download using hf and move to expected location
-hf download nvidia/NV-Segment-CT --local-dir models/ && \
-mv models/vista3d_pretrained_model/model.pt models/model.pt
 ```
 
+Model weights are prepared automatically during inference. The first run downloads the checkpoint from Hugging Face into the local Hugging Face cache and links it at `models/model.pt`; later runs reuse the cached weights while still touching Hugging Face download stats for each inference.
+
 ## 1.1 **NV-Segment-CT** [[Github]](https://github.com/NVIDIA-Medtech/NV-Segment-CTMR/tree/main/NV-Segment-CT) [[Huggingface]](https://huggingface.co/nvidia/NV-Segment-CT)
 
 ### Automatic Segmentation (support multi-gpu batch processing)
@@ -39,9 +36,9 @@ python -m monai.bundle run --config_file="['configs/inference.json', 'configs/ba
 # Automatic Batch segmentation for the whole folder with multi-gpu support. mgpu_inference.json is below. change nproc_per_node to your GPU number.
 torchrun --nproc_per_node=2 --nnodes=1 -m monai.bundle run --config_file="['configs/inference.json', 'configs/batch_inference.json', 'configs/mgpu_inference.json']" --input_dir="example/" --output_dir="example/"
 ```
-```
+
 Note: For more details about batch processing, please refer to NV-Segment-CTMR readme.md
-```
+
 ### Interactive segmentation
 
 ```bash
@@ -69,12 +66,14 @@ For more details, please refer to [this](inference.md).
 
 We provide predefined finetuning tutorial in [details](inference.md).
 For complicated finetuning, we suggest users to do vibe coding to generate finetuning pipelines by simply reuse the model and checkpoint
+
 ```python
 from monai.networks.nets.vista3d import vista3d132
 vista3d132.load_state_dict(pretrained_ckpt, strict=True)
 ```
 
 ## References
+
 - He, Yufan, et al. "VISTA3D: A unified segmentation foundation model for 3D medical imaging." Proceedings of the Computer Vision and Pattern Recognition Conference. 2025. <https://openaccess.thecvf.com/content/CVPR2025/html/He_VISTA3D_A_Unified_Segmentation_Foundation_Model_For_3D_Medical_Imaging_CVPR_2025_paper.html>
 
 ## License

NV-Segment-CT/docs/finetune.md

@@ -1,6 +1,6 @@
 # Finetune configurations
 
-### Step1: Generate Data json file
+## Step1: Generate Data json file
 
 Users need to provide a json data split for continuous learning (`configs/msd_task09_spleen_folds.json` from the [MSD](http://medicaldecathlon.com/) is provided as an example). The data split should meet the following format ('testing' labels are optional):
 
@@ -25,11 +25,11 @@ Example code for 5 fold cross-validation generation can be found [here](data.md)
 Note the data is not the absolute path to the image and label file. The actual image file will be `os.path.join(dataset_dir, data["training"][item]["image"])`, where `dataset_dir` is defined in `configs/train_continual.json`. Also 5-fold cross-validation is not required! `fold=0` is defined in train.json, which means any data item with fold==0 will be used as validation and other fold will be used for training. So if you only have train/val split, you can manually set validation data with "fold": 0 in its datalist and the other to be training by setting "fold" to any number other than 0.
 ```
 
-### Step2: Changing hyperparameters
+## Step2: Changing hyperparameters
 
 For continual learning, user can change `configs/train_continual.json`. More advanced users can change configurations in `configs/train.json`.  Most hyperparameters are straighforward and user can tell based on their names. The users must manually change the following keys in `configs/train_continual.json`.
 
-#### 1. `label_mappings`
+### 1. `label_mappings`
 
 ```json
     "label_mappings": {
@@ -54,15 +54,15 @@ For continual learning, user can change `configs/train_continual.json`. More adv
 If you cannot find a relevant semantic label for your class, just use any value < `num_classes` defined in train_continue.json.
 For more details about this label_mapping, please read [this](finetune.md).
 
-#### 2.  `data_list_file_path` and `dataset_dir`
+### 2.  `data_list_file_path` and `dataset_dir`
 
 Change `data_list_file_path` to the absolute path of your data json split. Change `dataset_dir` to the root folder that combines with the relative path in the data json split.
 
-#### 3. Optional hyperparameters and details are [here](finetune.md)
+### 3. Optional hyperparameters and details are [here](finetune.md)
 
 Hyperparameteers finetuning is important and varies from task to task.
 
-### Step3: Run finetuning
+## Step3: Run finetuning
 
 The hyperparameters in `configs/train_continual.json` will overwrite ones in `configs/train.json`. Configs in the back will overide the previous ones if they have the same key.
 
@@ -80,7 +80,7 @@ torchrun --nnodes=1 --nproc_per_node=8 -m monai.bundle run \
     --config_file="['configs/train.json','configs/train_continual.json','configs/multi_gpu_train.json']"
 ```
 
-#### MLFlow Visualization
+### MLFlow Visualization
 
 MLFlow is enabled by default (defined in train.json, use_mlflow) and the data is stored in the `mlruns/` folder under the bundle's root directory. To launch the MLflow UI and track your experiment data, follow these steps:
 
@@ -114,7 +114,7 @@ torchrun --nnodes=1 --nproc_per_node=8 -m monai.bundle run \
     --config_file="['configs/train.json','configs/train_continual.json','configs/evaluate.json','configs/mgpu_evaluate.json']"
 ```
 
-### Other explanatory items
+### Evaluation explanatory items
 
 The `label_mapping` in `evaluation.json` does not include `0` because the postprocessing step performs argmax (`VistaPostTransformd`), and a `0` prediction would negatively impact performance. In continuous learning, however, `0` is included for validation because no argmax is performed, and validation is done channel-wise (include_background=False). Additionally, `Relabeld` in `postprocessing` is required to map `label` and `pred` back to sequential indexes like `0, 1, 2, 3, 4` for dice calculation, as they are not in one-hot format. Evaluation does not support `point`, but finetuning does, as it does not perform argmax.
 
@@ -143,7 +143,6 @@ The `label_mapping` in `evaluation.json` does not include `0` because the postpr
 - Make sure you removed the `subclass` dictionary from inference.json if you ever mapped local index to [2,20,21]
 - Make sure `0` is not included in your inference prompt for automatic segmentation.
 
-
 ## Configurations
 
 ### Best practice to set label_mapping
@@ -183,7 +182,7 @@ In this bundle, the training is simplified by jointly training with class prompt
 NOTE: If user doesn't use interactive segmentation, set `drop_point_prob=1` and `drop_label_prob=0` in train.json might provide a faster and easier finetuning process.
 ```
 
-### Other explanatory items
+### Training explanatory items
 
 In `train.json`, `validate[evaluator][val_head]` can be `auto` and `point`. If `auto`, the validation results will be automatic segmentation. If `point`,
 the validation results will be sampling one positive point per object per patch. The validation scheme of combining auto and point is deprecated due to

NV-Segment-CT/requirements.txt

@@ -16,4 +16,4 @@ timm
 pytorch-ignite
 tensorboardX
 mlflow
-huggingface_hub
+huggingface_hub

NV-Segment-CT/scripts/__init__.py

@@ -12,7 +12,12 @@
 # from .evaluator import EnsembleEvaluator, Evaluator, SupervisedEvaluator
 # from .multi_gpu_supervised_trainer import create_multigpu_supervised_evaluator, create_multigpu_supervised_trainer
 
-from .early_stop_score_function import score_function
-
 # Ensures bundle expressions like ``scripts.batch_inference_utils.build_input_list`` resolve.
-from . import batch_inference_utils  # noqa: F401
+from . import batch_inference_utils as batch_inference_utils
+from .early_stop_score_function import score_function as score_function
+from .huggingface_download import (
+    prepare_huggingface_checkpoint as prepare_huggingface_checkpoint,
+)
+from .huggingface_download import (
+    touch_huggingface_download_counter as touch_huggingface_download_counter,
+)

NV-Segment-CT/scripts/batch_inference_utils.py

@@ -78,9 +78,7 @@ def collect_input_paths(
     for p in sorted(input_root.glob(pattern)):
         if not p.is_file():
             continue
-        if should_skip_path_by_parent_rules(
-            p, skip_dir_names=skip_dir_names, skip_dir_prefixes=skip_dir_prefixes
-        ):
+        if should_skip_path_by_parent_rules(p, skip_dir_names=skip_dir_names, skip_dir_prefixes=skip_dir_prefixes):
             continue
         out.append(p)
     return out
@@ -201,9 +199,7 @@ def build_input_list(
         )
         payload = {"paths": paths, "n_discovered": n_discovered}
         if use_cache and local_rank == "0":
-            cache = _cache_path(
-                input_dir, output_dir, output_postfix, output_ext, skip, names, prefixes
-            )
+            cache = _cache_path(input_dir, output_dir, output_postfix, output_ext, skip, names, prefixes)
             cache.parent.mkdir(parents=True, exist_ok=True)
             tmp = cache.with_suffix(".json.tmp")
             tmp.write_text(json.dumps(payload))
@@ -217,9 +213,7 @@ def build_input_list(
                 flush=True,
             )
     else:
-        cache = _cache_path(
-            input_dir, output_dir, output_postfix, output_ext, skip, names, prefixes
-        )
+        cache = _cache_path(input_dir, output_dir, output_postfix, output_ext, skip, names, prefixes)
         deadline = time.time() + wait_sec
         payload = None
         while time.time() < deadline:
@@ -228,9 +222,7 @@ def build_input_list(
                 break
             time.sleep(0.05)
         else:
-            raise RuntimeError(
-                "[nvseg] batch: cache timeout (raise batch_cache_wait_sec or set batch_use_input_list_cache false)"
-            )
+            raise RuntimeError("[nvseg] batch: cache timeout (raise batch_cache_wait_sec or set batch_use_input_list_cache false)")
         paths, n_discovered = _parse_cache_payload(payload)
         # Stale legacy `[]` or missing n_discovered: recompute so workers agree with rank 0.
         if n_discovered < 0:

NV-Segment-CT/scripts/huggingface_download.py

@@ -0,0 +1,85 @@
+import os
+import shutil
+from pathlib import Path
+
+
+def _is_rank_zero() -> bool:
+    for name in ("RANK", "LOCAL_RANK", "SLURM_PROCID"):
+        value = os.environ.get(name)
+        if value not in (None, "", "0"):
+            return False
+    return True
+
+
+def touch_huggingface_download_counter(
+    repo_id: str,
+    filename: str = "config.json",
+    revision: str = "main",
+    rank_zero_only: bool = True,
+) -> str | None:
+    """Force a tiny Hugging Face file request without re-downloading weights."""
+
+    if rank_zero_only and not _is_rank_zero():
+        return None
+
+    try:
+        from huggingface_hub import hf_hub_download
+    except ImportError:
+        print("[nvseg] warning: huggingface_hub is not installed; skipping Hugging Face download counter touch.")
+        return None
+
+    try:
+        path = hf_hub_download(
+            repo_id=repo_id,
+            filename=filename,
+            repo_type="model",
+            revision=revision,
+            force_download=True,
+        )
+    except Exception as exc:  # noqa: BLE001
+        print(f"[nvseg] warning: could not touch Hugging Face download counter for {repo_id}/{filename}: {exc}")
+        return None
+
+    print(f"[nvseg] touched Hugging Face download counter for {repo_id}/{filename}")
+    return path
+
+
+def prepare_huggingface_checkpoint(
+    repo_id: str,
+    checkpoint_filename: str,
+    local_checkpoint_path: str,
+    counter_filename: str = "config.json",
+    revision: str = "main",
+    rank_zero_only: bool = True,
+) -> str:
+    """Ensure the local MONAI checkpoint path exists and touch HF stats for this inference."""
+
+    local_path = Path(local_checkpoint_path)
+
+    if rank_zero_only and not _is_rank_zero():
+        return str(local_path)
+
+    touch_huggingface_download_counter(repo_id, counter_filename, revision, rank_zero_only=False)
+    if local_path.exists():
+        return str(local_path)
+
+    try:
+        from huggingface_hub import hf_hub_download
+    except ImportError as exc:
+        raise RuntimeError(f"{local_path} does not exist and huggingface_hub is not installed; cannot download {repo_id}.") from exc
+
+    checkpoint_path = hf_hub_download(
+        repo_id=repo_id,
+        filename=checkpoint_filename,
+        repo_type="model",
+        revision=revision,
+    )
+
+    local_path.parent.mkdir(parents=True, exist_ok=True)
+    try:
+        local_path.symlink_to(checkpoint_path)
+    except OSError:
+        shutil.copy2(checkpoint_path, local_path)
+
+    print(f"[nvseg] prepared checkpoint at {local_path}")
+    return str(local_path)

NV-Segment-CT/scripts/test_batch_inference.sh

@@ -0,0 +1,41 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+BUNDLE_ROOT=$(cd "${SCRIPT_DIR}/.." && pwd)
+PYTHON_BIN=${PYTHON_BIN:-python}
+WORK_DIR=${WORK_DIR:-$(mktemp -d)}
+KEEP_TEST_WORKDIR=${KEEP_TEST_WORKDIR:-0}
+
+if [[ "${KEEP_TEST_WORKDIR}" != "1" ]]; then
+    trap 'rm -rf "${WORK_DIR}"' EXIT
+fi
+
+INPUT_DIR="${WORK_DIR}/input"
+OUTPUT_DIR="${WORK_DIR}/output"
+mkdir -p "${INPUT_DIR}" "${OUTPUT_DIR}"
+
+cp "${BUNDLE_ROOT}/example/spleen_03.nii.gz" "${INPUT_DIR}/"
+
+cd "${BUNDLE_ROOT}"
+
+"${PYTHON_BIN}" -m monai.bundle run \
+    --config_file="['configs/inference.json', 'configs/batch_inference.json']" \
+    --input_dir="${INPUT_DIR}" \
+    --output_dir="${OUTPUT_DIR}" \
+    2>&1 | tee "${WORK_DIR}/batch_inference.log"
+
+EXPECTED_OUTPUT="${OUTPUT_DIR}/spleen_03/spleen_03_trans.nii.gz"
+test -s "${EXPECTED_OUTPUT}"
+grep -q "\[nvseg\] batch resume (skip existing outputs): 1 volume" "${WORK_DIR}/batch_inference.log"
+grep -q "\[nvseg\] touched Hugging Face download counter for nvidia/NV-Segment-CT/config.json" "${WORK_DIR}/batch_inference.log"
+
+"${PYTHON_BIN}" -m monai.bundle run \
+    --config_file="['configs/inference.json', 'configs/batch_inference.json']" \
+    --input_dir="${INPUT_DIR}" \
+    --output_dir="${OUTPUT_DIR}" \
+    2>&1 | tee "${WORK_DIR}/batch_resume.log"
+
+grep -q "\[nvseg\] batch: nothing to run (resume); ok" "${WORK_DIR}/batch_resume.log"
+
+echo "[nvseg-test] CT batch inference smoke test passed: ${EXPECTED_OUTPUT}"