nnunet-mlx

nnU-Net v2 inference in native MLX for Apple Silicon. Run existing nnU-Net models locally on Mac.

• Convert and run pretrained nnU-Net v2 model folders
• 2D and 3D PlainConvUNet and ResidualEncoderUNet
• FP32 using MLX; PyTorch is only needed to convert .pth checkpoints

Part of MedMLX. Technical notes · Compatibility

Quick start · Input and output · Technical details · License

Requirements

Apple Silicon Mac with Metal, uv, and conda or mamba (setup uses it to build a pinned Python 3.13 environment).

Quick start

Prepare the source checkout:

git clone https://github.com/MedMLX/nnunet-mlx.git
cd nnunet-mlx
make env                     # creates .venv with MLX on Metal
uv pip install --python .venv/bin/python -e '.[conversion]'   # to convert .pth checkpoints

Use .venv/bin/nnunet-mlx, or activate .venv. Every command accepts --help.

Convert an upstream nnU-Net model folder once:

nnunet-mlx convert model-folder \
  --source-model-folder /path/to/upstream-model \
  --output-model-folder mlx-model

Predict one case, passing one --input per modality in channel order:

nnunet-mlx model predict \
  --model-folder mlx-model \
  --input case_0000.nii.gz \
  --output-prefix predictions/case

Use nnunet-mlx model predict-folder for a whole folder. Add --overwrite to replace existing output, or --continue-existing to skip finished cases.

Only convert checkpoints you trust: .pth files are pickles. See SECURITY.md.

Python

from nnunet_mlx import ModelFolderPredictor, load_model_bundle

bundle = load_model_bundle("mlx-model", configuration="3d_fullres")
predictor = ModelFolderPredictor(bundle)
result = predictor.predict_case(
    ["case_0000.nii.gz"], "predictions/case", save_probabilities=False, overwrite=False,
)

For a single-channel 3D class model, predictor.predict_nifti(image) accepts a NiBabel NIfTI-1 image and returns local labels on its original grid without intermediate files. Fold logits stay on Metal except where restoration needs the host route (separate-Z or non-linear resampling). Cascades and regions use the file APIs above.

The model is also discoverable by medmlx as nnunet; see docs/technical.md for the runner.

Input and output

Supply one NIfTI file per input channel in the model's channel order. Prediction commands write segmentation files under the selected output prefix or directory; probability export is optional. Supported models and configurations are listed below and in the compatibility notes.

Anything else fails with a clear error rather than falling back. Pretrained-model support varies; see Compatibility.

Technical details

Targets nnU-Net v2.8.1 and dynamic-network-architectures v0.4.4. Networks, conversion and the prediction path are compared against the upstream PyTorch implementation on seeded-weight and synthetic inputs and on selected released models (see Compatibility); agreement on other models or data is not implied. Development, parity testing, the PancCTMultiTalentV2 staging recipe and the runner contract are in docs/technical.md; see also Architecture and Roadmap.

Intended use

Intended for research and education, not clinical diagnosis or treatment decisions. Outputs require human review.

License

Apache-2.0; see LICENSE and THIRD_PARTY_NOTICES.md. Kernels derived from MLX keep their MIT license. Pretrained nnU-Net weights have their own licenses and are not distributed here. Please cite via CITATION.cff and cite nnU-Net.

Implementation mapping

Central implementation map records upstream-to-native changes, source references and qualification scope.