Improved overall functionality and added functions for batch inference

README.md

@@ -14,6 +14,8 @@ library_name: transformers
 pipeline_tag: feature-extraction
 ---
 
+📢 [2026-07-17] SPECTRE now ships a command-line tool: point `spectre embed` at a `.nii`/`.nii.gz` file or a folder of them to get embeddings without writing any Python. Check below for details and usage examples.
+
 📢 [2026-05-20] The pretrained SPECTRE model can now be loaded directly through the `transformers` library, no separate SPECTRE package installation required. Check below for details and usage examples.
 
 📢 [2026-04-10] SPECTRE is now an official baseline for the [**CVPR 2026 Workshop Competition: Foundation Models for General CT Image Diagnosis**](https://www.codabench.org/competitions/12650/)! See `experiments/cvpr26_fm_for_ct_diag_task_1` for scripts and additional details.  
@@ -22,7 +24,6 @@ pipeline_tag: feature-extraction
 
 📢 [2026-01-20] [Semantic segmentation](https://github.com/cviviers/nnUNet) code and configurations using the nnUNet framework are now released!  
 
-
 # SPECTRE 👻👻👻
 
 <p align="center">
@@ -45,6 +46,7 @@ SPECTRE has been trained on a large cohort of **open-source CT scans** of the **
 This repository provides pretrained SPECTRE models together with tools for fine-tuning and evaluation.
 
 ## 🧠 Pretrained Models
+
 The pretrained SPECTRE model can easily be imported using the `transformers` library
 
 ```python
@@ -55,56 +57,75 @@ model = AutoModel.from_pretrained('cclaess/SPECTRE-Large', trust_remote_code=Tru
 or by using the `spectre-fm` package as follows:
 
 ```python
-from spectre import SpectreImageFeatureExtractor, MODEL_CONFIGS
-config = MODEL_CONFIGS['spectre-large-pretrained']
-model = SpectreImageFeatureExtractor.from_config(config)
+from spectre import SpectreImageFeatureExtractor
+model = SpectreImageFeatureExtractor.from_pretrained('spectre-large')
+```
+
+Run `spectre list-models` to see what is available. Pass `include_feature_combiner=False` for per-crop backbone features instead of one embedding per scan.
+
+### 🖥️ From the command line
+
+If you just want embeddings out of a CT scan and would rather not write Python:
+
+```bash
+pip install "spectre-fm[inference]"
+
+spectre embed scan.nii.gz -o embeddings/          # one scan
+spectre embed /data/scans/ -o embeddings/         # a whole folder
 ```
 
-A simple forward pass would look like:
+This handles everything internally with the defaults SPECTRE was pretrained on: HU windowing to [-1000, 1000], RAS orientation, and 128×128×64 crops at the scan's native voxel spacing. Each scan produces `<name>.npz` containing `cls` (one vector for the scan) and `patch_tokens` (one vector per crop, shaped to the crop grid), plus a `manifest.csv`. Useful flags: `--backbone-only`, `--device cuda`, `--spacing 0.5 0.5 1.0` to resample, and `--max-crops-per-forward` if you run out of memory. See `spectre embed --help`.
+
+### 🐍 From Python
+
+Hand the model a scan in Hounsfield Units and it does the windowing for you:
+
 ```python
 import torch
 
-model.eval()
+# One scan: (C, H, W, D) or (H, W, D), in raw HU.
+scan = torch.randn(1, 384, 384, 256) * 500 - 500
+
+with torch.no_grad():
+    features = model(scan)
+
+print("Features shape:", features.shape)   # (T', F') -> a CLS token plus one token per crop
+```
+
+Scans of different sizes can be embedded together by passing a list. All crops from all scans go through the backbone in a single pass, and only the feature combiner is split back out per scan:
+
+```python
+scans = [scan_a, scan_b, scan_c]           # any sizes, all in HU
 
-# Dummy input: (batch, crops, channels, height, width, depth)
-# For a (3 x 3 x 4) grid of (128 x 128 x 64) CT patches -> Total scan size (384 x 384 x 256)
-x = torch.randn(1, 1, 384, 384, 256)
-B, C, H, W, D = x.shape
+with torch.no_grad():
+    features = model.extract(scans)        # -> list of (T', F') tensors
+```
 
-patch_size = (128, 128, 64)
-pH, pW, pD = patch_size
+If you have already windowed the scans yourself, pass the crops and their grid instead, and they are used untouched:
+
+```python
+from spectre import window_scan
 
-x = x.view(
-  B, C,
-  H // pH, pH,
-  W // pW, pW,
-  D // pD, pD,
-).permute(0, 2, 4, 6, 1, 3, 5, 7).reshape(B, -1, C, pH, pW, pD)
+crops, grid_size = window_scan(scan)       # (N, C, 128, 128, 64), (n_h, n_w, n_d)
 
 with torch.no_grad():
-    features = model(
-      x, 
-      grid_size=(
-        H // pH,
-        W // pW,
-        D // pD,
-      ),
-    )
-print("Features shape:", features.shape)
+    features = model(crops, grid_size=grid_size)
 ```
 
+> **Reading files:** `spectre.load_ct` / `spectre.load_and_window` read `.nii`/`.nii.gz` and need the `[inference]` extra. Everything above works with just `pip install spectre-fm`.
+
 Alternatively, you can download the weights of the separate components through HuggingFace using the following links:
 
 | Architecture              | Input Modality     | Pretraining Objective   | Model Weights                                                                                                               |
 |---------------------------|--------------------|-------------------------|-----------------------------------------------------------------------------------------------------------------------------|
-| SPECTRE-ViT-Local         | CT crops           | SSL                     | [Link](https://huggingface.co/cclaess/SPECTRE/resolve/main/spectre_backbone_vit_large_patch16_128_no_vla.pt?download=true)  |
-| SPECTRE-ViT-Local         | CT crops           | SSL + VLA               | [Link](https://huggingface.co/cclaess/SPECTRE/resolve/main/spectre_backbone_vit_large_patch16_128.pt?download=true)         |
-| SPECTRE-ViT-Global        | Embedded CT crops  | VLA                     | [Link](https://huggingface.co/cclaess/SPECTRE/resolve/main/spectre_combiner_feature_vit_large.pt?download=true)             |
-| Qwen3-Embedding-0.6B LoRA | Text (radiology)   | VLA                     | [Link](https://huggingface.co/cclaess/SPECTRE/resolve/main/spectre_qwen3_embedding_0.6B_lora.pt?download=true)              |
+| SPECTRE-ViT-Local         | CT crops           | SSL                     | [Download](https://huggingface.co/cclaess/SPECTRE/resolve/main/spectre_backbone_vit_large_patch16_128_no_vla.pt?download=true)  |
+| SPECTRE-ViT-Local         | CT crops           | SSL + VLA               | [Download](https://huggingface.co/cclaess/SPECTRE/resolve/main/spectre_backbone_vit_large_patch16_128.pt?download=true)         |
+| SPECTRE-ViT-Global        | Embedded CT crops  | VLA                     | [Download](https://huggingface.co/cclaess/SPECTRE/resolve/main/spectre_combiner_feature_vit_large.pt?download=true)             |
+| Qwen3-Embedding-0.6B LoRA | Text (radiology)   | VLA                     | [Download](https://huggingface.co/cclaess/SPECTRE/resolve/main/spectre_qwen3_embedding_0.6B_lora.pt?download=true)              |
 
 ## 🩻 Segmentation (nnUNet)
 
-If you're looking for a nnUNet-based segmentation pipeline that uses SPECTRE as the backbone, see: https://github.com/cviviers/nnUNet
+If you're looking for a nnUNet-based segmentation pipeline that uses SPECTRE as the backbone, see this [GitHub](https://github.com/cviviers/nnUNet).
 
 ## 📂 Repository Contents
 
@@ -131,6 +152,12 @@ pip install spectre-fm
 
 This installs only the runtime dependencies needed to load and run the pretrained models.
 
+To read CT scans from `.nii`/`.nii.gz` files or use the `spectre` command-line tool, add the inference extra:
+
+```bash
+pip install "spectre-fm[inference]"
+```
+
 If you want to fine-tune or pretrain SPECTRE, install the matching extra:
 
 ```bash
@@ -168,7 +195,9 @@ pip install git+https://github.com/cclaess/SPECTRE.git
 To facilitate deployment and reproducibility, SPECTRE can be run using **Docker**. This allows you to set up a fully functional environment without manually installing dependencies using your own local copy of spectre.
 
 ### **Building the Docker Image**
+
 First, ensure you have **Docker** installed. Then, clone and navigate to the repository to build the image:
+
 ```bash
 git clone https://github.com/cclaess/SPECTRE
 cd SPECTRE
@@ -176,22 +205,28 @@ docker build -t spectre-fm .
 ```
 
 ### **Running Experiments Inside Docker**
+
 Once the image is built, you can start a container and execute scripts inside it. For example, to run a DINO pretraining experiment:
+
 ```bash
 docker run --gpus all --rm -v "$(pwd):/mnt" spectre-fm python3 experiments/pretraining/pretrain_dino.py --config_file spectre/configs/dino_default.yaml --output_dir /mnt/outputs/pretraining/dino/
 ```
+
 - `--gpus all` enables GPU acceleration if available.
 - `--rm` removes the container after execution.
 - `-v $(pwd):/mnt` mounts the current directory inside the container.
 
 ## ⚖️ License
+
 - **Code: MIT** — see `LICENSE` (permissive; commercial use permitted).
 - **Pretrained model weights: CC-BY-NC-SA** — non-commercial share-alike. The weights and any derivative models that include these weights are NOT cleared for commercial use. See `LICENSE_MODELS` for details and the precise license text.
 
 > Note: the pretrained weights are subject to the original dataset licenses. Users intending to use SPECTRE in commercial settings should verify dataset and model licensing and obtain any required permissions.
 
 ## 📜 Citation
+
 If you use SPECTRE in your research or wish to cite it, please use the following BibTeX entry of our [preprint](https://arxiv.org/abs/2511.17209):
+
 ```
 @misc{claessens_scaling_2025,
   title = {Scaling {Self}-{Supervised} and {Cross}-{Modal} {Pretraining} for {Volumetric} {CT} {Transformers}},
@@ -203,6 +238,7 @@ If you use SPECTRE in your research or wish to cite it, please use the following
 ```
 
 ## 🤝 Acknowledgements
+
 This project builds upon prior work in self-supervised learning, medical imaging, and transformer-based representation learning. We especially acknowledge [**MONAI**](https://project-monai.github.io/) for their awesome framework and the [**timm**](https://timm.fast.ai/) & [**lightly**](https://docs.lightly.ai/self-supervised-learning/) Python libraries for providing 2D PyTorch models (timm) and object-oriented self-supervised learning methods (lightly), from which we adapted parts of the code for 3D.
 
-[![Star History Chart](https://api.star-history.com/svg?repos=cclaess/SPECTRE&type=Date)](https://star-history.com/#cclaess/SPECTRE&Date)
+[![Star History Chart](https://api.star-history.com/svg?repos=cclaess/SPECTRE&type=Date)](https://star-history.com/#cclaess/SPECTRE&Date)

config.json

@@ -16,6 +16,11 @@
     }
   },
   "backbone_name": "vit_large_patch16_128",
+  "crop_size": [
+    128,
+    128,
+    64
+  ],
   "dtype": "float32",
   "feature_combiner_kwargs": {
     "global_pool": "",
@@ -27,6 +32,14 @@
     }
   },
   "feature_combiner_name": "feat_vit_large",
+  "hu_range": [
+    -1000.0,
+    1000.0
+  ],
+  "include_feature_combiner": true,
   "model_type": "spectre",
-  "transformers_version": "5.3.0"
+  "orientation": "RAS",
+  "preset": "spectre-large",
+  "transformers_version": "5.3.0",
+  "voxel_spacing": null
 }

configuration_spectre.py

@@ -1,32 +1,69 @@
 from transformers import PretrainedConfig
 
+from spectre.presets import get_preset
+from spectre.windowing import DEFAULT_CROP_SIZE, DEFAULT_HU_RANGE
+
 
 class SpectreConfig(PretrainedConfig):
+    """Configuration of a pretrained SPECTRE model.
+
+    Nothing about the model is written down twice. The architecture is resolved from a named
+    preset in `spectre.presets`, and the preprocessing defaults come from `spectre.windowing` -
+    the same values the package itself uses. The resolved architecture is still serialised into
+    `config.json`, so the published config stays self-describing.
+
+    The fields above the architecture are the input contract: SPECTRE consumes a grid of
+    `crop_size` crops cut from a CT scan oriented to `orientation` and windowed from `hu_range`
+    onto [0, 1] at `voxel_spacing` (None meaning the scan's native spacing).
+    `SpectreImageFeatureExtractor` applies all of that for you when handed a raw scan in
+    Hounsfield Units.
+
+    Compatibility runs both ways:
+    * A `config.json` written before the `preset` field existed carries an explicit
+      `backbone_name` and `backbone_kwargs`; those take precedence, so it keeps loading its own
+      architecture instead of silently inheriting this class's defaults.
+    * Older code reading a newer `config.json` ignores the fields it does not know and reads the
+      resolved architecture, which is still written out in full.
+    """
+
     model_type = "spectre"
 
     def __init__(
         self,
-        backbone_name="vit_large_patch16_128",
-        backbone_kwargs={
-            "num_classes": 0,
-            "global_pool": '',
-            "pos_embed": "rope",
-            "rope_kwargs": {"base": 1000.0},
-            "init_values": 1.0,
-        },
-        feature_combiner_name="feat_vit_large",
-        feature_combiner_kwargs={
-            "num_classes": 0,
-            "global_pool": "",
-            "pos_embed": "rope",
-            "rope_kwargs": {"base": 100.0},
-            "init_values": 1.0,
-        },
+        preset="spectre-large",
+        # What the model expects to be fed.
+        crop_size=DEFAULT_CROP_SIZE,
+        hu_range=DEFAULT_HU_RANGE,
+        voxel_spacing=None,
+        orientation="RAS",
+        # Architecture. Left as None these are resolved from `preset`; passing them explicitly
+        # (as every pre-0.2 config.json does) takes precedence.
+        backbone_name=None,
+        backbone_kwargs=None,
+        feature_combiner_name=None,
+        feature_combiner_kwargs=None,
+        include_feature_combiner=True,
         **kwargs,
     ):
         super().__init__(**kwargs)
 
+        self.preset = preset
+        self.crop_size = tuple(crop_size)
+        self.hu_range = tuple(hu_range)
+        self.voxel_spacing = tuple(voxel_spacing) if voxel_spacing is not None else None
+        self.orientation = orientation
+        self.include_feature_combiner = include_feature_combiner
+
+        if backbone_name is None:
+            resolved = get_preset(preset)
+            backbone_name = resolved.backbone
+            if backbone_kwargs is None:
+                backbone_kwargs = resolved.backbone_kwargs
+            feature_combiner_name = feature_combiner_name or resolved.feature_combiner
+            if feature_combiner_kwargs is None:
+                feature_combiner_kwargs = resolved.feature_combiner_kwargs
+
         self.backbone_name = backbone_name
         self.backbone_kwargs = backbone_kwargs or {}
-        self.feature_combiner_name = feature_combiner_name
-        self.feature_combiner_kwargs = feature_combiner_kwargs or {}
+        self.feature_combiner_name = feature_combiner_name if include_feature_combiner else None
+        self.feature_combiner_kwargs = feature_combiner_kwargs or {}

modeling_spectre.py

@@ -3,6 +3,7 @@ from transformers import PreTrainedModel
 from transformers.modeling_outputs import BaseModelOutput
 
 from spectre.model import SpectreImageFeatureExtractor
+
 try:
     from .configuration_spectre import SpectreConfig
 except ImportError:
@@ -10,6 +11,21 @@ except ImportError:
 
 
 class SpectreModel(PreTrainedModel):
+    """`transformers` wrapper around `SpectreImageFeatureExtractor`.
+
+        from transformers import AutoModel
+        model = AutoModel.from_pretrained("cclaess/SPECTRE-Large", trust_remote_code=True)
+
+        # A raw CT scan (C, H, W, D) in Hounsfield Units - windowed internally.
+        features = model(scan).last_hidden_state
+
+    This wrapper takes a single tensor, because `BaseModelOutput` holds a tensor. For batches of
+    differently-sized scans - where all crops share one backbone pass - use the underlying
+    extractor directly, which accepts a list:
+
+        features = model.model.extract([scan_a, scan_b])   # -> list of (T', F')
+    """
+
     config_class = SpectreConfig
     base_model_prefix = "spectre"
     main_input_name = "pixel_values"
@@ -26,6 +42,11 @@ class SpectreModel(PreTrainedModel):
 
         self.post_init()
 
+    @property
+    def crop_size(self):
+        """Spatial size (H, W, D) of one CT crop the backbone consumes."""
+        return self.model.crop_size
+
     def forward(
         self,
         pixel_values: torch.Tensor,
@@ -33,9 +54,12 @@ class SpectreModel(PreTrainedModel):
         return_dict=False,
         **kwargs,
     ):
-        outputs = self.model(pixel_values, grid_size=grid_size)
+        """Embed one or more CT scans. See `SpectreImageFeatureExtractor.forward` for the shapes
+        accepted; `grid_size` is required only for input that is already windowed into crops.
+        """
+        outputs = self.model(pixel_values, grid_size=grid_size, **kwargs)
 
         if not return_dict:
             return outputs
 
-        return BaseModelOutput(last_hidden_state=outputs)
+        return BaseModelOutput(last_hidden_state=outputs)

spectre/__init__.py

@@ -2,25 +2,22 @@
 
 Expose a small, stable public API here so users can do:
 
-	from spectre import SpectreImageFeatureExtractor, models
+	from spectre import SpectreImageFeatureExtractor
+	model = SpectreImageFeatureExtractor.from_pretrained("spectre-large")
 
 Keep implementations in subpackages; this file only re-exports the most
 important symbols and subpackages for convenience.
+
+Note for maintainers: `scripts/export_hf.py` rewrites this file when vendoring the package into
+the Hugging Face Hub repo, dropping the imports it does not ship. Never import `.cli` here, or
+anything that needs an optional dependency at module level.
 """
 from .model import SpectreImageFeatureExtractor, MODEL_CONFIGS
+from .presets import SpectrePreset, PRESETS, get_preset, list_presets, list_pretrained
+from .windowing import window_scan, largest_multiple_center_crop, grid_patch, scale_intensity_range
 from . import models
 from . import utils
-
-__version__ = "0.1.0"
-__author__ = "Cris Claessens"
-__email__ = "c.h.b.claessens@tue.nl"
-
-__all__ = [
-	"SpectreImageFeatureExtractor",
-	"MODEL_CONFIGS",
-	"models",
-	"utils",
-	"__version__",
-    "__author__",
-    "__email__",
-]
+__version__ = '0.2.1'
+__author__ = 'Cris Claessens'
+__email__ = 'c.h.b.claessens@tue.nl'
+__all__ = ['SpectreImageFeatureExtractor', 'MODEL_CONFIGS', 'SpectrePreset', 'PRESETS', 'get_preset', 'list_presets', 'list_pretrained', 'window_scan', 'largest_multiple_center_crop', 'grid_patch', 'scale_intensity_range', 'models', 'utils', '__version__', '__author__', '__email__']

spectre/model.py

@@ -1,141 +1,149 @@
+import copy
+import warnings
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
+
 import torch
 import torch.nn as nn
 
+from spectre.presets import (
+    BACKBONE_KWARGS,
+    COMBINER_KWARGS,
+    get_preset,
+    list_presets,
+    list_pretrained,
+)
+from spectre.windowing import window_scan
+
+
+def _legacy_model_configs() -> Dict[str, Dict[str, Any]]:
+    """The pre-0.2 `MODEL_CONFIGS` payload.
 
-MODEL_CONFIGS = {
-    "spectre-small": {
-        "name": "spectre-small",
-        "backbone": "vit_small_patch16_128",
-        "backbone_checkpoint_path_or_url": None,
-        "backbone_kwargs": {
-            "num_classes": 0,
-            "global_pool": '',
-            "pos_embed": "rope",
-            "rope_kwargs": {"base": 1000.0},
-            "init_values": 1.0,
-        },
-        "feature_combiner": "feat_vit_small",
-        "feature_combiner_checkpoint_path_or_url": None,
-        "feature_combiner_kwargs": {
-            "num_classes": 0,
-            "global_pool": '',
-            "pos_embed": "rope",
-            "rope_kwargs": {"base": 100.0},
-            "init_values": 1.0,
-        },
-        "description": "SPECTRE model with ViT-Small backbone and feature combiner.",
-    },  # Pretrained/Distilled checkpoints will be added later
-    "spectre-base": {
-        "name": "spectre-base",
-        "backbone": "vit_base_patch16_128",
-        "backbone_checkpoint_path_or_url": None,
-        "backbone_kwargs": {
-            "num_classes": 0,
-            "global_pool": '',
-            "pos_embed": "rope",
-            "rope_kwargs": {"base": 1000.0},
-            "init_values": 1.0,
-        },
-        "feature_combiner": "feat_vit_base",
-        "feature_combiner_checkpoint_path_or_url": None,
-        "feature_combiner_kwargs": {
-            "num_classes": 0,
-            "global_pool": '',
-            "pos_embed": "rope",
-            "rope_kwargs": {"base": 100.0},
-            "init_values": 1.0,
-        },
-        "description": "SPECTRE model with ViT-Base backbone and feature combiner.",
-    },  # Pretrained/Distilled checkpoints will be added later
-    "spectre-large": {
-        "name": "spectre-large",
-        "backbone": "vit_large_patch16_128",
-        "backbone_checkpoint_path_or_url": None,
-        "backbone_kwargs": {
-            "num_classes": 0,
-            "global_pool": '',
-            "pos_embed": "rope",
-            "rope_kwargs": {"base": 1000.0},
-            "init_values": 1.0,
-        },
-        "feature_combiner": "feat_vit_large",
-        "feature_combiner_checkpoint_path_or_url": None,
-        "feature_combiner_kwargs": {
-            "num_classes": 0,
-            "global_pool": '',
-            "pos_embed": "rope",
-            "rope_kwargs": {"base": 100.0},
-            "init_values": 1.0,
-        },
-        "description": "SPECTRE model with ViT-Large backbone and feature combiner.",
-    },
-    "spectre-large-pretrained": {
-        "name": "spectre-large-pretrained",
-        "backbone": "vit_large_patch16_128",
-        "backbone_checkpoint_path_or_url": "https://huggingface.co/cclaess/SPECTRE/resolve/main/spectre_backbone_vit_large_patch16_128.pt?download=true",
-        "backbone_kwargs": {
-            "num_classes": 0,
-            "global_pool": '',
-            "pos_embed": "rope",
-            "rope_kwargs": {"base": 1000.0},
-            "init_values": 1.0,
-        },
-        "feature_combiner": "feat_vit_large",
-        "feature_combiner_checkpoint_path_or_url": "https://huggingface.co/cclaess/SPECTRE/resolve/main/spectre_combiner_feature_vit_large.pt?download=true",
-        "feature_combiner_kwargs": {
-            "num_classes": 0,
-            "global_pool": '',
-            "pos_embed": "rope",
-            "rope_kwargs": {"base": 100.0},
-            "init_values": 1.0,
-        },
-        "description": "Pretrained SPECTRE model with ViT-Large backbone and feature combiner.",
+    The architecture kwargs come from `spectre.presets`, but the *layout* is reproduced by hand
+    because the two registries disagree about what `"spectre-large"` means: here it is a
+    weightless architecture and `"spectre-large-pretrained"` carries the weights, whereas
+    `from_pretrained("spectre-large")` loads weights. Deriving the mapping would silently change
+    which checkpoint existing callers get. `tests/test_legacy_api.py` pins the result.
+    """
+    def entry(name, backbone, combiner, backbone_ckpt, combiner_ckpt, description):
+        return {
+            "name": name,
+            "backbone": backbone,
+            "backbone_checkpoint_path_or_url": backbone_ckpt,
+            "backbone_kwargs": copy.deepcopy(BACKBONE_KWARGS),
+            "feature_combiner": combiner,
+            "feature_combiner_checkpoint_path_or_url": combiner_ckpt,
+            "feature_combiner_kwargs": copy.deepcopy(COMBINER_KWARGS),
+            "description": description,
+        }
+
+    return {
+        "spectre-small": entry(
+            "spectre-small", "vit_small_patch16_128", "feat_vit_small", None, None,
+            "SPECTRE model with ViT-Small backbone and feature combiner.",
+        ),
+        "spectre-base": entry(
+            "spectre-base", "vit_base_patch16_128", "feat_vit_base", None, None,
+            "SPECTRE model with ViT-Base backbone and feature combiner.",
+        ),
+        "spectre-large": entry(
+            "spectre-large", "vit_large_patch16_128", "feat_vit_large", None, None,
+            "SPECTRE model with ViT-Large backbone and feature combiner.",
+        ),
+        "spectre-large-pretrained": entry(
+            "spectre-large-pretrained", "vit_large_patch16_128", "feat_vit_large",
+            "https://huggingface.co/cclaess/SPECTRE/resolve/main/spectre_backbone_vit_large_patch16_128.pt?download=true",
+            "https://huggingface.co/cclaess/SPECTRE/resolve/main/spectre_combiner_feature_vit_large.pt?download=true",
+            "Pretrained SPECTRE model with ViT-Large backbone and feature combiner.",
+        ),
     }
-}
+
+
+_MODEL_CONFIGS_DEPRECATION = (
+    "spectre.MODEL_CONFIGS is deprecated since spectre-fm 0.2.1 and will be removed in 0.3.0. "
+    "Use SpectreImageFeatureExtractor.from_pretrained('spectre-large') for the pretrained model, "
+    "or spectre.list_presets() to see the available architectures."
+)
+
+
+class _DeprecatedModelConfigs(dict):
+    """`MODEL_CONFIGS` that warns on access.
+
+    A dict subclass rather than a module-level `__getattr__`, because `spectre/__init__.py`
+    imports the name eagerly - a module hook would fire on every `import spectre`, including
+    inside the Hub's trust_remote_code import.
+    """
+
+    def __getitem__(self, key):
+        warnings.warn(_MODEL_CONFIGS_DEPRECATION, FutureWarning, stacklevel=2)
+        return super().__getitem__(key)
+
+    def get(self, key, default=None):
+        # dict.get does not route through __getitem__ in CPython, so it needs its own warning.
+        warnings.warn(_MODEL_CONFIGS_DEPRECATION, FutureWarning, stacklevel=2)
+        return super().get(key, default)
+
+
+MODEL_CONFIGS = _DeprecatedModelConfigs(_legacy_model_configs())
 
 
 class SpectreImageFeatureExtractor(nn.Module):
+    """SPECTRE image encoder: a ViT backbone over CT crops, plus an optional feature combiner.
+
+    The backbone embeds each `crop_size` crop independently; the feature combiner then attends
+    over the grid of crop embeddings to produce one embedding for the whole scan.
+
+    The usual way in is `from_pretrained`:
+
+        model = SpectreImageFeatureExtractor.from_pretrained("spectre-large")
+        features = model(scan)                       # scan: (C, H, W, D) in Hounsfield Units
+        features = model([scan_a, scan_b])           # differing sizes, one backbone pass
+
+    Pass `include_feature_combiner=False` for per-crop backbone features only.
+    """
+
     def __init__(
-        self, 
-        backbone_name: str, 
-        backbone_kwargs: dict = {},
+        self,
+        backbone_name: str,
+        backbone_kwargs: Optional[dict] = None,
         backbone_checkpoint_path_or_url: str | None = None,
-        feature_combiner_name: str | None = None, 
-        feature_combiner_kwargs: dict = {},
+        feature_combiner_name: str | None = None,
+        feature_combiner_kwargs: Optional[dict] = None,
         feature_combiner_checkpoint_path_or_url: str | None = None,
         **kwargs,
-    ):  
+    ):
+        # This signature is a published contract: the Hugging Face wrapper
+        # (hf_export/modeling_spectre.py) calls it with exactly these keywords.
         super().__init__()
         self.backbone = None
         self.feature_combiner = None
         self._init_backbone(
             backbone_name,
             checkpoint_path_or_url=backbone_checkpoint_path_or_url,
-            **backbone_kwargs,
+            **(backbone_kwargs or {}),
             **kwargs,
         )
         if feature_combiner_name is not None:
             self._init_feature_combiner(
                 feature_combiner_name,
                 checkpoint_path_or_url=feature_combiner_checkpoint_path_or_url,
-                **feature_combiner_kwargs,
+                **(feature_combiner_kwargs or {}),
                 **kwargs,
             )
 
     def _init_backbone(
-        self, 
+        self,
         model_name: str,
         checkpoint_path_or_url: str | None = None,
         **kwargs
     ):
         backbone_cls = getattr(__import__('spectre.models', fromlist=[model_name]), model_name)
         self.backbone = backbone_cls(
-            checkpoint_path_or_url=checkpoint_path_or_url, 
+            checkpoint_path_or_url=checkpoint_path_or_url,
             **kwargs,
         )
-        
+
     def _init_feature_combiner(
-        self, 
+        self,
         model_name: str,
         checkpoint_path_or_url: str | None = None,
         **kwargs,
@@ -147,20 +155,105 @@ class SpectreImageFeatureExtractor(nn.Module):
 
         feature_combiner_cls = getattr(__import__('spectre.models', fromlist=[model_name]), model_name)
         self.feature_combiner = feature_combiner_cls(
-            patch_dim=patch_dim, 
-            checkpoint_path_or_url=checkpoint_path_or_url, 
+            patch_dim=patch_dim,
+            checkpoint_path_or_url=checkpoint_path_or_url,
             **kwargs,
         )
 
+    @classmethod
+    def from_pretrained(
+        cls,
+        name: str = "spectre-large",
+        *,
+        include_feature_combiner: bool = True,
+        pretrained: bool = True,
+        strict: bool = True,
+        backbone_kwargs: Optional[dict] = None,
+        feature_combiner_kwargs: Optional[dict] = None,
+        device: Union[str, torch.device, None] = None,
+        dtype: Optional[torch.dtype] = None,
+        verbose: bool = False,
+    ) -> 'SpectreImageFeatureExtractor':
+        """Build a named SPECTRE model, by default with its published weights.
+
+        Args:
+            name: A preset name, e.g. "spectre-large". See `spectre.list_presets()`.
+            include_feature_combiner: Include the feature combiner (one embedding per scan). When
+                False you get per-crop backbone features and no combiner checkpoint is fetched.
+            pretrained: Load published weights. Raises if the preset has none.
+            strict: Require the checkpoint to match the architecture exactly.
+            backbone_kwargs: Overrides merged over the preset's backbone kwargs.
+            feature_combiner_kwargs: Overrides merged over the preset's combiner kwargs.
+            device: Device to move the model to.
+            dtype: Dtype to cast the model to.
+            verbose: Print checkpoint download/load progress.
+
+        Returns:
+            The model in eval mode.
+        """
+        if any(token in str(name) for token in ('/', '\\', '://')) or str(name).endswith('.pt'):
+            raise ValueError(
+                f"from_pretrained() takes a preset name, not a path or URL, but got {name!r}. "
+                f"Available presets: {list_presets()}. To load your own checkpoint, construct "
+                f"the model directly:\n"
+                f"    SpectreImageFeatureExtractor(\n"
+                f"        backbone_name='vit_large_patch16_128',\n"
+                f"        backbone_checkpoint_path_or_url={name!r},\n"
+                f"    )"
+            )
+
+        preset = get_preset(name)
+
+        if pretrained and not preset.has_pretrained_weights:
+            raise ValueError(
+                f"Preset {name!r} has no published weights. Presets with weights: "
+                f"{list_pretrained()}. Pass pretrained=False for a randomly initialised "
+                f"{name!r} model."
+            )
+
+        merged_backbone_kwargs = dict(preset.backbone_kwargs)
+        merged_backbone_kwargs.update(backbone_kwargs or {})
+        merged_combiner_kwargs = dict(preset.feature_combiner_kwargs)
+        merged_combiner_kwargs.update(feature_combiner_kwargs or {})
+
+        # __init__ forwards **kwargs into both sub-model factories, which only accept verbose and
+        # strict when they are actually loading a checkpoint.
+        load_kwargs = {"verbose": verbose, "strict": strict} if pretrained else {}
+
+        model = cls(
+            backbone_name=preset.backbone,
+            backbone_kwargs=merged_backbone_kwargs,
+            backbone_checkpoint_path_or_url=preset.backbone_weights if pretrained else None,
+            feature_combiner_name=preset.feature_combiner if include_feature_combiner else None,
+            feature_combiner_kwargs=merged_combiner_kwargs,
+            feature_combiner_checkpoint_path_or_url=(
+                preset.feature_combiner_weights if pretrained and include_feature_combiner else None
+            ),
+            **load_kwargs,
+        )
+
+        if device is not None or dtype is not None:
+            model = model.to(device=device, dtype=dtype)
+        return model.eval()
+
+    @property
+    def crop_size(self) -> Tuple[int, int, int]:
+        """Spatial size (H, W, D) of one CT crop the backbone consumes."""
+        return tuple(self.backbone.patch_embed.img_size)
+
+    @property
+    def has_feature_combiner(self) -> bool:
+        return self.feature_combiner is not None
+
     def extract_backbone_features(
-        self, 
+        self,
         x: torch.Tensor,
     ):
         """
-        Extract features from the backbone for a batch of image sets. Input is expected to be of 
-        shape (B, N, C, H, W, D), where B is the batch size, N is the number of image patches per 
+        Extract features from the backbone for a batch of image sets. Input is expected to be of
+        shape (B, N, C, H, W, D), where B is the batch size, N is the number of image patches per
         image, C is the number of channels, H is height, W is width, and D is depth.
-        The output will be a tensor of extracted features (B, N, T, F) where T is the number of 
+        The output will be a tensor of extracted features (B, N, T, F) where T is the number of
         tokens and F is the feature dimension.
 
         Args:
@@ -178,7 +271,7 @@ class SpectreImageFeatureExtractor(nn.Module):
         return features
 
     def combine_features(
-        self, 
+        self,
         features: torch.Tensor,
         grid_size: tuple[int, int, int],
     ):
@@ -207,21 +300,244 @@ class SpectreImageFeatureExtractor(nn.Module):
         features = self.feature_combiner(features, grid_size)  # (B, T', F')
         return features
 
-    def forward(self, x, grid_size: tuple[int, int, int] | None = None):
-        features = self.extract_backbone_features(x)
-        if self.feature_combiner is not None:
-            assert grid_size is not None, \
-                "`grid_size` must be provided when using feature combiner"
-            features = self.combine_features(features, grid_size)
-        return features
+    def _pool_crop_tokens(self, features: torch.Tensor) -> torch.Tensor:
+        """Reduce per-crop backbone tokens (M, T, F) to the combiner's input (M, 2F).
+
+        Must stay bit-identical to `combine_features`: the same CLS token concatenated with the
+        same mean over the same axis. The prefix index 0 assumes num_prefix_tokens == 1, which
+        holds for every preset; changing it would move published embeddings.
+        """
+        if features.ndim == 2:  # backbone already pooled globally
+            return features
+        if features.shape[1] == 1:  # CLS only
+            return features[:, 0]
+        return torch.cat([features[:, 0], features[:, 1:].mean(dim=1)], dim=-1)
+
+    def _encode(
+        self,
+        items: List[Tuple[torch.Tensor, Optional[Tuple[int, int, int]]]],
+        max_crops_per_forward: Optional[int] = None,
+    ) -> List[torch.Tensor]:
+        """Run crops from any number of scans through the backbone in one pass.
+
+        `items` pairs each scan's crops (N_i, C, cH, cW, cD) with its grid. Crops from every scan
+        go through the backbone together; only the feature combiner, whose positional encoding
+        depends on the grid, is split back out per scan.
+        """
+        if not items:
+            return []
+
+        crops_per_scan = [crops for crops, _ in items]
+        grids = [grid for _, grid in items]
+
+        reference = crops_per_scan[0].shape[1:]
+        for i, crops in enumerate(crops_per_scan):
+            if crops.shape[1:] != reference:
+                raise ValueError(
+                    f"All scans must have the same crop shape, but scan 0 has crops of "
+                    f"{tuple(reference)} and scan {i} has {tuple(crops.shape[1:])}."
+                )
+
+        param = next(self.parameters())
+        counts = [int(crops.shape[0]) for crops in crops_per_scan]
+        all_crops = torch.cat(
+            [crops.to(device=param.device, dtype=param.dtype) for crops in crops_per_scan], dim=0
+        )
+
+        chunk = max_crops_per_forward if max_crops_per_forward else all_crops.shape[0]
+        chunk_features = []
+        for crops in torch.split(all_crops, chunk, dim=0):
+            features = self.backbone(crops)
+            if features.ndim == 2:  # only CLS token
+                features = features.unsqueeze(1)
+            if self.feature_combiner is not None:
+                # Pool inside the loop: unpooled tokens for a whole batch of scans are orders of
+                # magnitude larger than the pooled features the combiner actually needs.
+                features = self._pool_crop_tokens(features)
+            chunk_features.append(features)
+
+        features = torch.cat(chunk_features, dim=0)
+        per_scan = torch.split(features, counts, dim=0)
+
+        if self.feature_combiner is None:
+            return list(per_scan)
+
+        # The combiner's rotary embedding depends on the grid and there is no attention mask, so
+        # scans can only share a forward pass when their grids match exactly.
+        groups: Dict[Tuple[int, int, int], List[int]] = {}
+        for index, grid in enumerate(grids):
+            groups.setdefault(tuple(grid), []).append(index)
+
+        results: List[torch.Tensor] = [torch.empty(0)] * len(items)
+        for grid, indices in groups.items():
+            batch = torch.stack([per_scan[i] for i in indices], dim=0)  # (G, N, 2F)
+            combined = self.feature_combiner(batch, grid)  # (G, T', F')
+            for position, index in enumerate(indices):
+                results[index] = combined[position]
+
+        return results
+
+    def _prepare_scan(
+        self,
+        x: torch.Tensor,
+        grid_size: Optional[Sequence[int]],
+        index: Optional[int] = None,
+    ) -> Tuple[torch.Tensor, Tuple[int, int, int]]:
+        """Normalise one scan to (crops, grid), windowing it first if it is still a whole volume."""
+        where = "" if index is None else f" (scan {index})"
+
+        if hasattr(x, 'as_tensor'):  # MONAI MetaTensor, without importing monai
+            x = x.as_tensor()
+        if not isinstance(x, torch.Tensor):
+            raise TypeError(f"Expected a torch.Tensor{where}, got {type(x).__name__}.")
+
+        crop_size = self.crop_size
+
+        if x.ndim in (3, 4):
+            if grid_size is not None:
+                raise ValueError(
+                    f"grid_size must not be passed for a raw {x.ndim}D scan{where}: it is derived "
+                    f"from the scan's shape while windowing. Pass grid_size only for "
+                    f"pre-windowed (N, C, H, W, D) input."
+                )
+            expected_channels = self.backbone.patch_embed.proj.in_channels
+            if x.ndim == 4 and x.shape[0] != expected_channels:
+                raise ValueError(
+                    f"A 4D input{where} is read as one raw scan (C, H, W, D) with C="
+                    f"{expected_channels}, but got shape {tuple(x.shape)}. For a batch of scans, "
+                    f"pass a list of tensors instead."
+                )
+            return window_scan(x, crop_size)
+
+        if x.ndim == 5:
+            # (N, C, H, W, D) and (B, C, H, W, D) are indistinguishable by rank, so 5D always
+            # means one pre-windowed scan; raw batches go through a list.
+            if tuple(x.shape[2:]) != crop_size:
+                raise ValueError(
+                    f"A 5D input{where} is read as one pre-windowed scan (N, C, H, W, D) whose "
+                    f"crops must be {crop_size}, but got crops of {tuple(x.shape[2:])}. If this "
+                    f"is a batch of raw scans, pass a list of (C, H, W, D) tensors instead."
+                )
+            if grid_size is None:
+                raise ValueError(
+                    f"grid_size is required for pre-windowed (N, C, H, W, D) input{where}. Pass "
+                    f"the raw scan as (C, H, W, D) to have it windowed for you."
+                )
+            grid = tuple(int(g) for g in grid_size)
+            if len(grid) != 3:
+                raise ValueError(f"grid_size must have 3 elements (H, W, D), got {grid}.")
+            expected = grid[0] * grid[1] * grid[2]
+            if x.shape[0] != expected:
+                raise ValueError(
+                    f"grid_size {grid} implies {expected} crops but got {x.shape[0]}{where}."
+                )
+            return x, grid
+
+        raise ValueError(
+            f"Unsupported input with {x.ndim} dimensions{where}. Expected one of: (H, W, D) or "
+            f"(C, H, W, D) for a raw scan, (N, C, H, W, D) for a pre-windowed scan, "
+            f"(B, N, C, H, W, D) for a pre-windowed batch, or a list of scans."
+        )
+
+    def forward(
+        self,
+        x: Union[torch.Tensor, Sequence[torch.Tensor]],
+        grid_size: Optional[Union[Sequence[int], Sequence[Sequence[int]]]] = None,
+        *,
+        max_crops_per_forward: Optional[int] = None,
+    ) -> Union[torch.Tensor, List[torch.Tensor]]:
+        """Embed one or more CT scans.
+
+        Accepts, dispatching on rank:
+
+        | input                | meaning                        | grid_size | returns          |
+        |----------------------|--------------------------------|-----------|------------------|
+        | `(H, W, D)`          | raw scan, in HU                | derived   | `(T', F')`       |
+        | `(C, H, W, D)`       | raw scan, in HU                | derived   | `(T', F')`       |
+        | `(N, C, cH, cW, cD)` | pre-windowed scan              | required  | `(T', F')`       |
+        | `(B, N, C, ...)`     | pre-windowed batch, one grid   | required  | `(B, T', F')`    |
+        | `list[...]`          | batch of scans, may differ     | per scan  | `list[(T', F')]` |
+
+        Raw input is HU-windowed and cut into crops internally; pre-windowed input is passed
+        through untouched, so it must already be preprocessed the same way.
+
+        Without a feature combiner the trailing shape is the backbone's per-crop tokens
+        `(N, T, F)` instead of `(T', F')`.
+
+        Args:
+            x: Scan(s), as above.
+            grid_size: Crop grid for pre-windowed input. For a list, either one grid shared by
+                every scan or one per scan.
+            max_crops_per_forward: Cap the crops in a single backbone pass to bound memory.
+                None runs them all at once.
+        """
+        is_batch = isinstance(x, (list, tuple))
+
+        if is_batch:
+            grids = self._grids_for_batch(grid_size, len(x))
+            items = [self._prepare_scan(scan, g, i) for i, (scan, g) in enumerate(zip(x, grids))]
+            return self._encode(items, max_crops_per_forward)
+
+        if hasattr(x, 'as_tensor'):
+            x = x.as_tensor()
+
+        if isinstance(x, torch.Tensor) and x.ndim == 6:
+            if grid_size is None:
+                raise ValueError(
+                    "grid_size is required for pre-windowed (B, N, C, H, W, D) input. Pass a list "
+                    "of raw (C, H, W, D) scans to have them windowed for you."
+                )
+            grid = tuple(int(g) for g in grid_size)
+            items = [self._prepare_scan(scan, grid, i) for i, scan in enumerate(x)]
+            return torch.stack(self._encode(items, max_crops_per_forward), dim=0)
+
+        items = [self._prepare_scan(x, grid_size)]
+        return self._encode(items, max_crops_per_forward)[0]
+
+    @staticmethod
+    def _grids_for_batch(
+        grid_size: Optional[Union[Sequence[int], Sequence[Sequence[int]]]],
+        count: int,
+    ) -> List[Optional[Tuple[int, ...]]]:
+        """Resolve `grid_size` for a list input: one shared grid, one per scan, or none."""
+        if grid_size is None:
+            return [None] * count
+        # A grid is always three ints, so all-ints means one shared grid.
+        if all(isinstance(g, int) for g in grid_size):
+            return [tuple(int(g) for g in grid_size)] * count
+        grids = [tuple(int(v) for v in g) for g in grid_size]
+        if len(grids) != count:
+            raise ValueError(
+                f"Got {len(grids)} grid sizes for {count} scans. Pass one grid_size per scan, or "
+                f"a single (H, W, D) grid shared by all of them."
+            )
+        return grids
+
+    @torch.inference_mode()
+    def extract(
+        self,
+        scans: Union[torch.Tensor, Sequence[torch.Tensor]],
+        grid_size: Optional[Union[Sequence[int], Sequence[Sequence[int]]]] = None,
+        *,
+        max_crops_per_forward: Optional[int] = 16,
+    ) -> Union[torch.Tensor, List[torch.Tensor]]:
+        """`forward` without gradients, with memory capped by default. See `forward`."""
+        return self.forward(scans, grid_size, max_crops_per_forward=max_crops_per_forward)
 
     @classmethod
     def from_config(
-        cls, 
+        cls,
         config: dict,
         **kwargs,
     ) -> 'SpectreImageFeatureExtractor':
-        
+        warnings.warn(
+            "SpectreImageFeatureExtractor.from_config() is deprecated since spectre-fm 0.2.1 and "
+            "will be removed in 0.3.0. Use "
+            "SpectreImageFeatureExtractor.from_pretrained('spectre-large') instead, or call "
+            "SpectreImageFeatureExtractor(...) directly for a custom architecture.",
+            FutureWarning,
+            stacklevel=2,
+        )
         model = cls(
             backbone_name=config["backbone"],
             backbone_checkpoint_path_or_url=config.get("backbone_checkpoint_path_or_url", None),

spectre/models/__init__.py

@@ -1,31 +1,3 @@
-from .vision_transformer import (
-    VisionTransformer,
-    vit_tiny_patch16_128, 
-    vit_small_patch16_128, 
-    vit_base_patch16_128, 
-    vit_base_patch32_128,
-    vit_large_patch16_128,
-    vit_large_patch32_128,
-)
-from .vision_transformer_features import (
-    FeatureVisionTransformer,
-    feat_vit_tiny,
-    feat_vit_small,
-    feat_vit_base,
-    feat_vit_large,
-)
-
-__all__ = [
-    'VisionTransformer',
-    'vit_tiny_patch16_128',
-    'vit_small_patch16_128',
-    'vit_base_patch16_128',
-    'vit_base_patch32_128',
-    'vit_large_patch16_128',
-    'vit_large_patch32_128',
-    'FeatureVisionTransformer',
-    'feat_vit_tiny',
-    'feat_vit_small',
-    'feat_vit_base',
-    'feat_vit_large',
-]
+from .vision_transformer import VisionTransformer, vit_tiny_patch16_128, vit_small_patch16_128, vit_base_patch16_128, vit_base_patch32_128, vit_large_patch16_128, vit_large_patch32_128
+from .vision_transformer_features import FeatureVisionTransformer, feat_vit_tiny, feat_vit_small, feat_vit_base, feat_vit_large
+__all__ = ['VisionTransformer', 'vit_tiny_patch16_128', 'vit_small_patch16_128', 'vit_base_patch16_128', 'vit_base_patch32_128', 'vit_large_patch16_128', 'vit_large_patch32_128', 'FeatureVisionTransformer', 'feat_vit_tiny', 'feat_vit_small', 'feat_vit_base', 'feat_vit_large']

spectre/models/vision_transformer.py

@@ -1,6 +1,5 @@
 import os
 from functools import partial
-from urllib.parse import urlparse
 from typing import (
     Tuple, Union, Callable, Literal, 
     Optional, Type, Set, List, Dict, Any,
@@ -10,17 +9,17 @@ import torch
 import torch.nn as nn
 from timm.layers import PatchDropout, AttentionPoolLatent
 from timm.models.vision_transformer import LayerScale, DropPath, Mlp
-from huggingface_hub import hf_hub_download, load_state_dict_from_file
 
 from spectre.models.layers import (
-    PatchEmbed, 
-    Attention, 
+    PatchEmbed,
+    Attention,
     RotaryPositionEmbedding,
 )
 from spectre.utils import (
-    resample_abs_pos_embed, 
-    feature_take_indices, 
+    resample_abs_pos_embed,
+    feature_take_indices,
     global_pool_nlc,
+    load_pretrained_into,
 )
 
 
@@ -591,53 +590,13 @@ class VisionTransformer(nn.Module):
             cls,
             checkpoint_path_or_url: Union[str, os.PathLike],
             verbose: bool = True,
+            strict: bool = True,
             **kwargs
     ) -> 'VisionTransformer':
         """Load pretrained model weights from a local path or a URL."""
         model = cls(**kwargs)
-
-        def _is_url(path: str) -> bool:
-            try:
-                parsed = urlparse(str(path))
-                return parsed.scheme in ('http', 'https')
-            except Exception:
-                return False
-            
-        def _is_hf_url(path: str) -> bool:
-            try:
-                parsed = urlparse(str(path))
-                return 'huggingface.co' in parsed.netloc
-            except Exception:
-                return False
-
-        if _is_hf_url(checkpoint_path_or_url):
-            if verbose:
-                print(f"Downloading pretrained weights from Hugging Face URL: {checkpoint_path_or_url}")
-            # Extract repo_id and filename from the URL
-            parsed = urlparse(checkpoint_path_or_url)
-            parts = parsed.path.strip('/').split('/')
-            repo_id = '/'.join(parts[:2])  # e.g., 'cclaess/SPECTRE'
-            filename = parts[-1]           # e.g., 'spectre_backbone_vit_large_patch16_128.pt'
-
-            local_path = hf_hub_download(repo_id=repo_id, filename=filename)
-            state_dict = load_state_dict_from_file(local_path, map_location='cpu')
-        elif _is_url(checkpoint_path_or_url):
-            if verbose:
-                print(f"Downloading pretrained weights from URL: {checkpoint_path_or_url}")
-            state_dict = torch.hub.load_state_dict_from_url(
-                checkpoint_path_or_url, map_location='cpu', weights_only=False, progress=verbose)
-        else:
-            local_path = os.fspath(checkpoint_path_or_url)
-            if not os.path.exists(local_path):
-                raise FileNotFoundError(f"Checkpoint file not found: {local_path}")
-        if verbose:
-            print(f"Loading checkpoint from local path: {local_path}")
-            state_dict = torch.load(local_path, map_location='cpu', weights_only=False)
-
-        msg = model.load_state_dict(state_dict, strict=False)
-        if verbose:
-            print(f"Loaded pretrained weights with msg: {msg}")
-        return model
+        return load_pretrained_into(
+            model, checkpoint_path_or_url, strict=strict, verbose=verbose)
 
 
 def vit_tiny_patch16_128(

spectre/models/vision_transformer_features.py

@@ -1,16 +1,19 @@
 import os
 import math
 from functools import partial
-from urllib.parse import urlparse
 from typing import Union, Callable, Literal, Optional, Type, Set, Tuple
 
 import torch
 import torch.nn as nn
 from timm.models.vision_transformer import Mlp
 from timm.layers import PatchDropout, AttentionPoolLatent
-from huggingface_hub import hf_hub_download, load_state_dict_from_file
 
-from spectre.utils import  global_pool_nlc, to_3tuple, resample_abs_pos_embed
+from spectre.utils import (
+    global_pool_nlc,
+    to_3tuple,
+    resample_abs_pos_embed,
+    load_pretrained_into,
+)
 from spectre.models.vision_transformer import Block
 from spectre.models.layers import RotaryPositionEmbedding
 
@@ -124,6 +127,11 @@ class FeatureVisionTransformer(nn.Module):
                 num_heads=num_heads,
                 **rope_kwargs,
             )
+            self.requires_per_sample_rope = any([
+                self.rope.shift_coords is not None,
+                self.rope.jitter_coords is not None,
+                self.rope.rescale_coords is not None,
+            ])
         self.pos_drop = nn.Dropout(p=pos_drop_rate)
         if patch_drop_rate > 0:
             self.patch_drop = PatchDropout(
@@ -318,53 +326,13 @@ class FeatureVisionTransformer(nn.Module):
             cls,
             checkpoint_path_or_url: Union[str, os.PathLike],
             verbose: bool = True,
+            strict: bool = True,
             **kwargs
     ) -> 'FeatureVisionTransformer':
         """Load pretrained model weights from a local path or a URL."""
         model = cls(**kwargs)
-
-        def _is_url(path: str) -> bool:
-            try:
-                parsed = urlparse(str(path))
-                return parsed.scheme in ('http', 'https')
-            except Exception:
-                return False
-            
-        def _is_hf_url(path: str) -> bool:
-            try:
-                parsed = urlparse(str(path))
-                return 'huggingface.co' in parsed.netloc
-            except Exception:
-                return False
-
-        if _is_hf_url(checkpoint_path_or_url):
-            if verbose:
-                print(f"Downloading pretrained weights from Hugging Face URL: {checkpoint_path_or_url}")
-            # Extract repo_id and filename from the URL
-            parsed = urlparse(checkpoint_path_or_url)
-            parts = parsed.path.strip('/').split('/')
-            repo_id = '/'.join(parts[:2])  # e.g., 'cclaess/SPECTRE'
-            filename = parts[-1]           # e.g., 'spectre_backbone_vit_large_patch16_128.pt'
-
-            local_path = hf_hub_download(repo_id=repo_id, filename=filename)
-            state_dict = load_state_dict_from_file(local_path, map_location='cpu')
-        elif _is_url(checkpoint_path_or_url):
-            if verbose:
-                print(f"Downloading pretrained weights from URL: {checkpoint_path_or_url}")
-            state_dict = torch.hub.load_state_dict_from_url(
-                checkpoint_path_or_url, map_location='cpu', weights_only=False, progress=verbose)
-        else:
-            local_path = os.fspath(checkpoint_path_or_url)
-            if not os.path.exists(local_path):
-                raise FileNotFoundError(f"Checkpoint file not found: {local_path}")
-        if verbose:
-            print(f"Loading checkpoint from local path: {local_path}")
-            state_dict = torch.load(local_path, map_location='cpu', weights_only=False)
-
-        msg = model.load_state_dict(state_dict, strict=False)
-        if verbose:
-            print(f"Loaded pretrained weights with msg: {msg}")
-        return model
+        return load_pretrained_into(
+            model, checkpoint_path_or_url, strict=strict, verbose=verbose)
 
 
 def feat_vit_tiny(

spectre/presets.py

@@ -0,0 +1,130 @@
+"""Named SPECTRE architectures and their published weights.
+
+This module is vendored into the Hugging Face Hub repository by `scripts/export_hf.py`, so it
+must import nothing beyond the standard library.
+"""
+from __future__ import annotations
+
+import copy
+import difflib
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+__all__ = [
+    'SpectrePreset',
+    'PRESETS',
+    'BACKBONE_KWARGS',
+    'COMBINER_KWARGS',
+    'get_preset',
+    'list_presets',
+    'list_pretrained',
+]
+
+_HF_WEIGHTS = "https://huggingface.co/cclaess/SPECTRE/resolve/main/{}?download=true"
+
+# The only place the SPECTRE architecture kwargs are written down. `spectre.model`, the Hugging
+# Face SpectreConfig and scripts/export_hf.py all resolve back to here rather than restating them.
+# The backbone sees raw CT crops; the feature combiner sees the backbone's per-crop embeddings,
+# hence the different RoPE bases.
+BACKBONE_KWARGS: Dict[str, Any] = {
+    "num_classes": 0,
+    "global_pool": '',
+    "pos_embed": "rope",
+    "rope_kwargs": {"base": 1000.0},
+    "init_values": 1.0,
+}
+COMBINER_KWARGS: Dict[str, Any] = {
+    "num_classes": 0,
+    "global_pool": '',
+    "pos_embed": "rope",
+    "rope_kwargs": {"base": 100.0},
+    "init_values": 1.0,
+}
+
+
+@dataclass(frozen=True)
+class SpectrePreset:
+    """One named SPECTRE architecture, optionally with published weights.
+
+    There is deliberately no crop size here: it is `backbone.patch_embed.img_size`, so the
+    architecture stays the single source of truth.
+    """
+    name: str
+    backbone: str
+    feature_combiner: str
+    backbone_kwargs: Dict[str, Any] = field(default_factory=dict)
+    feature_combiner_kwargs: Dict[str, Any] = field(default_factory=dict)
+    backbone_weights: Optional[str] = None
+    feature_combiner_weights: Optional[str] = None
+    description: str = ""
+
+    @property
+    def has_pretrained_weights(self) -> bool:
+        return self.backbone_weights is not None
+
+
+PRESETS: Dict[str, SpectrePreset] = {
+    "spectre-small": SpectrePreset(
+        name="spectre-small",
+        backbone="vit_small_patch16_128",
+        feature_combiner="feat_vit_small",
+        backbone_kwargs=copy.deepcopy(BACKBONE_KWARGS),
+        feature_combiner_kwargs=copy.deepcopy(COMBINER_KWARGS),
+        description="SPECTRE with a ViT-Small backbone. No published weights yet.",
+    ),
+    "spectre-base": SpectrePreset(
+        name="spectre-base",
+        backbone="vit_base_patch16_128",
+        feature_combiner="feat_vit_base",
+        backbone_kwargs=copy.deepcopy(BACKBONE_KWARGS),
+        feature_combiner_kwargs=copy.deepcopy(COMBINER_KWARGS),
+        description="SPECTRE with a ViT-Base backbone. No published weights yet.",
+    ),
+    "spectre-large": SpectrePreset(
+        name="spectre-large",
+        backbone="vit_large_patch16_128",
+        feature_combiner="feat_vit_large",
+        backbone_kwargs=copy.deepcopy(BACKBONE_KWARGS),
+        feature_combiner_kwargs=copy.deepcopy(COMBINER_KWARGS),
+        backbone_weights=_HF_WEIGHTS.format("spectre_backbone_vit_large_patch16_128.pt"),
+        feature_combiner_weights=_HF_WEIGHTS.format("spectre_combiner_feature_vit_large.pt"),
+        description="SPECTRE with a ViT-Large backbone, pretrained with SSL + vision-language alignment.",
+    ),
+}
+
+
+def list_presets() -> List[str]:
+    """Every known architecture name."""
+    return list(PRESETS)
+
+
+def list_pretrained() -> List[str]:
+    """The architecture names that have published weights."""
+    return [name for name, preset in PRESETS.items() if preset.has_pretrained_weights]
+
+
+def get_preset(name: str) -> SpectrePreset:
+    """Look up a preset by name.
+
+    The returned preset owns deep copies of its kwargs, so callers can mutate them without
+    corrupting the registry for the rest of the process.
+    """
+    try:
+        preset = PRESETS[name]
+    except KeyError:
+        suggestion = difflib.get_close_matches(str(name), list(PRESETS), n=1)
+        hint = f" Did you mean {suggestion[0]!r}?" if suggestion else ""
+        raise ValueError(
+            f"Unknown SPECTRE preset {name!r}. Available: {list_presets()}.{hint}"
+        ) from None
+
+    return SpectrePreset(
+        name=preset.name,
+        backbone=preset.backbone,
+        feature_combiner=preset.feature_combiner,
+        backbone_kwargs=copy.deepcopy(preset.backbone_kwargs),
+        feature_combiner_kwargs=copy.deepcopy(preset.feature_combiner_kwargs),
+        backbone_weights=preset.backbone_weights,
+        feature_combiner_weights=preset.feature_combiner_weights,
+        description=preset.description,
+    )

spectre/utils/__init__.py

@@ -1,65 +1,3 @@
-from ._utils import (
-    fix_random_seeds, 
-    to_ntuple,
-    to_1tuple,
-    to_2tuple,
-    to_3tuple,
-    to_4tuple,
-)
-from .modeling import (
-    deactivate_requires_grad_and_to_eval,
-    activate_requires_grad_and_to_train,
-    update_momentum,
-    update_drop_path_rate,
-    repeat_token,
-    expand_index_like,
-    get_at_index,
-    set_at_index,
-    mask_at_index,
-    mask_bool,
-    patchify,
-    random_token_mask,
-    resample_abs_pos_embed,
-    resample_abs_pos_embed_nhwdc,
-    resample_patch_embed,
-    feature_take_indices,
-    global_pool_nlc,
-    cat_keep_shapes,
-    uncat_with_shapes,
-    last_token_pool,
-    Format,
-    nchwd_to,
-    nhwdc_to,
-)
-
-__all__ = [
-    "fix_random_seeds",
-    "to_ntuple",
-    "to_1tuple",
-    "to_2tuple",
-    "to_3tuple",
-    "to_4tuple",
-    "deactivate_requires_grad_and_to_eval",
-    "activate_requires_grad_and_to_train",
-    "update_momentum",
-    "update_drop_path_rate",
-    "repeat_token",
-    "expand_index_like",
-    "get_at_index",
-    "set_at_index",
-    "mask_at_index",
-    "mask_bool",
-    "patchify",
-    "random_token_mask",
-    "resample_abs_pos_embed",
-    "resample_abs_pos_embed_nhwdc",
-    "resample_patch_embed",
-    "feature_take_indices",
-    "global_pool_nlc",
-    "cat_keep_shapes",
-    "uncat_with_shapes",
-    "last_token_pool",
-    "Format",
-    "nchwd_to",
-    "nhwdc_to",
-]
+from ._utils import fix_random_seeds, to_ntuple, to_1tuple, to_2tuple, to_3tuple, to_4tuple
+from .modeling import load_checkpoint, load_pretrained_into, deactivate_requires_grad_and_to_eval, activate_requires_grad_and_to_train, update_momentum, update_drop_path_rate, repeat_token, expand_index_like, get_at_index, set_at_index, mask_at_index, mask_bool, patchify, random_token_mask, resample_abs_pos_embed, resample_abs_pos_embed_nhwdc, resample_patch_embed, feature_take_indices, global_pool_nlc, cat_keep_shapes, uncat_with_shapes, last_token_pool, Format, nchwd_to, nhwdc_to
+__all__ = ['fix_random_seeds', 'to_ntuple', 'to_1tuple', 'to_2tuple', 'to_3tuple', 'to_4tuple', 'load_checkpoint', 'load_pretrained_into', 'deactivate_requires_grad_and_to_eval', 'activate_requires_grad_and_to_train', 'update_momentum', 'update_drop_path_rate', 'repeat_token', 'expand_index_like', 'get_at_index', 'set_at_index', 'mask_at_index', 'mask_bool', 'patchify', 'random_token_mask', 'resample_abs_pos_embed', 'resample_abs_pos_embed_nhwdc', 'resample_patch_embed', 'feature_take_indices', 'global_pool_nlc', 'cat_keep_shapes', 'uncat_with_shapes', 'last_token_pool', 'Format', 'nchwd_to', 'nhwdc_to']

spectre/utils/modeling.py

@@ -1,13 +1,130 @@
 from __future__ import annotations
 
+import os
 import math
+import warnings
 from enum import Enum
-from typing import List, Tuple, Optional, Union
+from typing import Dict, List, Tuple, Optional, Union
+from urllib.parse import urlparse
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import numpy as np
+from huggingface_hub import hf_hub_download, load_state_dict_from_file
+
+
+def _is_url(path: Union[str, os.PathLike]) -> bool:
+    try:
+        return urlparse(str(path)).scheme in ('http', 'https')
+    except Exception:
+        return False
+
+
+def _is_hf_url(path: Union[str, os.PathLike]) -> bool:
+    try:
+        return 'huggingface.co' in urlparse(str(path)).netloc
+    except Exception:
+        return False
+
+
+def _torch_load(path: Union[str, os.PathLike], map_location) -> Dict[str, torch.Tensor]:
+    """Load a checkpoint, preferring the safe unpickler and falling back if the file needs it."""
+    try:
+        return torch.load(path, map_location=map_location, weights_only=True)
+    except Exception:
+        warnings.warn(
+            f"Could not load {path} with weights_only=True, retrying with weights_only=False. "
+            f"This executes arbitrary code from the checkpoint - only do this for files you trust.",
+            RuntimeWarning,
+            stacklevel=2,
+        )
+        return torch.load(path, map_location=map_location, weights_only=False)
+
+
+def load_checkpoint(
+    checkpoint_path_or_url: Union[str, os.PathLike],
+    map_location: Union[str, torch.device] = 'cpu',
+    verbose: bool = False,
+) -> Dict[str, torch.Tensor]:
+    """Load a state dict from a local path, a Hugging Face URL, or a plain URL.
+
+    Args:
+        checkpoint_path_or_url: Local file, a `https://huggingface.co/<repo>/resolve/<rev>/<file>`
+            URL, or any other http(s) URL.
+        map_location: Passed through to the underlying loader.
+        verbose: Print where the checkpoint is coming from.
+    Returns:
+        The loaded state dict.
+    """
+    if _is_hf_url(checkpoint_path_or_url):
+        if verbose:
+            print(f"Downloading pretrained weights from Hugging Face URL: {checkpoint_path_or_url}")
+        parsed = urlparse(str(checkpoint_path_or_url))
+        parts = parsed.path.strip('/').split('/')
+        if len(parts) < 2:
+            raise ValueError(
+                f"Cannot parse a Hugging Face repo id out of {checkpoint_path_or_url!r}. "
+                f"Expected https://huggingface.co/<owner>/<repo>/resolve/<revision>/<filename>."
+            )
+        repo_id = '/'.join(parts[:2])
+        filename = parts[-1]
+        local_path = hf_hub_download(repo_id=repo_id, filename=filename)
+        return load_state_dict_from_file(local_path, map_location=map_location)
+
+    if _is_url(checkpoint_path_or_url):
+        if verbose:
+            print(f"Downloading pretrained weights from URL: {checkpoint_path_or_url}")
+        try:
+            return torch.hub.load_state_dict_from_url(
+                str(checkpoint_path_or_url), map_location=map_location,
+                weights_only=True, progress=verbose,
+            )
+        except Exception:
+            warnings.warn(
+                f"Could not load {checkpoint_path_or_url} with weights_only=True, retrying with "
+                f"weights_only=False. This executes arbitrary code from the checkpoint - only do "
+                f"this for URLs you trust.",
+                RuntimeWarning,
+                stacklevel=2,
+            )
+            return torch.hub.load_state_dict_from_url(
+                str(checkpoint_path_or_url), map_location=map_location,
+                weights_only=False, progress=verbose,
+            )
+
+    local_path = os.fspath(checkpoint_path_or_url)
+    if not os.path.exists(local_path):
+        raise FileNotFoundError(f"Checkpoint file not found: {local_path}")
+    if verbose:
+        print(f"Loading checkpoint from local path: {local_path}")
+    return _torch_load(local_path, map_location)
+
+
+def load_pretrained_into(
+    model: nn.Module,
+    checkpoint_path_or_url: Union[str, os.PathLike],
+    strict: bool = True,
+    verbose: bool = False,
+) -> nn.Module:
+    """Load `checkpoint_path_or_url` into `model` in place, with an actionable error on mismatch.
+
+    With `strict=False` a mismatched architecture silently leaves layers randomly initialised,
+    so the default here is `strict=True`.
+    """
+    state_dict = load_checkpoint(checkpoint_path_or_url, map_location='cpu', verbose=verbose)
+    try:
+        msg = model.load_state_dict(state_dict, strict=strict)
+    except RuntimeError as e:
+        raise RuntimeError(
+            f"Checkpoint {checkpoint_path_or_url} does not match {type(model).__name__}. This "
+            f"usually means the architecture kwargs differ from the ones the checkpoint was "
+            f"trained with. Pass strict=False to load the matching subset anyway (the rest stays "
+            f"randomly initialised).\n\n{e}"
+        ) from e
+    if verbose:
+        print(f"Loaded pretrained weights with msg: {msg}")
+    return model
 
 
 def deactivate_requires_grad_and_to_eval(model: nn.Module):

spectre/windowing.py

@@ -0,0 +1,245 @@
+"""Pure-torch CT preprocessing: HU scaling, centre cropping and grid patching.
+
+This module is vendored into the Hugging Face Hub repository by `scripts/export_hf.py`,
+so it must import nothing beyond `torch` and the standard library.
+
+The functions here reproduce the MONAI pipeline used for pretraining and evaluation
+(`ScaleIntensityRanged` -> `LargestMultipleCenterCropd` -> `GridPatchd`) exactly. Two details
+are load-bearing and are covered by `tests/test_windowing_parity.py`:
+
+* Patch order is C-order with depth varying fastest, i.e. crop `n` is at grid position
+  `n = (h * n_w + w) * n_d + d`. This matches both MONAI's `GridPatch` and the coordinates
+  `RotaryPositionEmbedding` builds, so getting it wrong shifts every embedding silently.
+* The centre-crop offset is `start = (S // 2) - (R // 2)`, which is not the same as
+  `(S - R) // 2` when the crop is odd and the volume even.
+"""
+from __future__ import annotations
+
+import warnings
+from typing import Optional, Sequence, Tuple
+
+import torch
+import torch.nn.functional as F
+
+DEFAULT_CROP_SIZE: Tuple[int, int, int] = (128, 128, 64)  # (H, W, D)
+DEFAULT_HU_RANGE: Tuple[float, float] = (-1000.0, 1000.0)
+
+__all__ = [
+    'DEFAULT_CROP_SIZE',
+    'DEFAULT_HU_RANGE',
+    'scale_intensity_range',
+    'largest_multiple_crop_size',
+    'largest_multiple_center_crop',
+    'grid_patch',
+    'window_scan',
+]
+
+
+def scale_intensity_range(
+    x: torch.Tensor,
+    a_min: float = DEFAULT_HU_RANGE[0],
+    a_max: float = DEFAULT_HU_RANGE[1],
+    b_min: float = 0.0,
+    b_max: float = 1.0,
+    clip: bool = True,
+) -> torch.Tensor:
+    """Map intensities from `[a_min, a_max]` onto `[b_min, b_max]`.
+
+    Mirrors `monai.transforms.ScaleIntensityRange`. SPECTRE was pretrained with the default
+    HU window of [-1000, 1000] -> [0, 1] with clipping.
+    """
+    if a_max == a_min:
+        raise ValueError(f"a_min and a_max must differ, got a_min == a_max == {a_min}.")
+    x = (x - a_min) / (a_max - a_min)
+    x = x * (b_max - b_min) + b_min
+    if clip:
+        x = torch.clamp(x, min(b_min, b_max), max(b_min, b_max))
+    return x
+
+
+def largest_multiple_crop_size(
+    spatial_shape: Sequence[int],
+    crop_size: Sequence[int] = DEFAULT_CROP_SIZE,
+) -> Tuple[int, ...]:
+    """Largest per-axis size that is a whole multiple of `crop_size`.
+
+    An axis shorter than its crop keeps its original size (there is no non-zero multiple to
+    take), matching `spectre.transforms.LargestMultipleCenterCropd`.
+    """
+    if len(spatial_shape) != len(crop_size):
+        raise ValueError(
+            f"spatial_shape {tuple(spatial_shape)} and crop_size {tuple(crop_size)} must have "
+            f"the same number of dimensions."
+        )
+    return tuple(
+        (s // c) * c if s >= c else s
+        for s, c in zip(spatial_shape, crop_size)
+    )
+
+
+def largest_multiple_center_crop(
+    x: torch.Tensor,
+    crop_size: Sequence[int] = DEFAULT_CROP_SIZE,
+) -> torch.Tensor:
+    """Centre-crop a channel-first `(C, H, W, D)` volume to a whole multiple of `crop_size`."""
+    if x.ndim != 4:
+        raise ValueError(f"Expected a channel-first (C, H, W, D) volume, got shape {tuple(x.shape)}.")
+
+    spatial_shape = tuple(x.shape[1:])
+    roi_size = largest_multiple_crop_size(spatial_shape, crop_size)
+
+    # MONAI's CenterSpatialCrop offset: roi_center - (roi_size // 2), clamped at 0. This is not
+    # the same as (S - R) // 2 for odd R and even S, so keep the formula verbatim.
+    slices = []
+    for size, roi in zip(spatial_shape, roi_size):
+        start = max(size // 2 - roi // 2, 0)
+        slices.append(slice(start, start + roi))
+    return x[(slice(None), *slices)]
+
+
+def grid_patch(
+    x: torch.Tensor,
+    crop_size: Sequence[int] = DEFAULT_CROP_SIZE,
+) -> Tuple[torch.Tensor, Tuple[int, int, int]]:
+    """Tile a `(C, H, W, D)` volume into non-overlapping crops.
+
+    Every spatial axis must already be a whole multiple of `crop_size` (see
+    `largest_multiple_center_crop`).
+
+    Returns:
+        The crops `(N, C, cH, cW, cD)` and the grid `(n_h, n_w, n_d)` with `N == n_h*n_w*n_d`.
+        Crop `n` sits at grid position `n = (h * n_w + w) * n_d + d`.
+    """
+    if x.ndim != 4:
+        raise ValueError(f"Expected a channel-first (C, H, W, D) volume, got shape {tuple(x.shape)}.")
+
+    channels = x.shape[0]
+    spatial_shape = tuple(x.shape[1:])
+    c_h, c_w, c_d = (int(c) for c in crop_size)
+
+    for axis, (size, crop) in enumerate(zip(spatial_shape, (c_h, c_w, c_d))):
+        if crop <= 0:
+            raise ValueError(f"crop_size must be positive, got {tuple(crop_size)}.")
+        if size % crop != 0:
+            raise ValueError(
+                f"Axis {'HWD'[axis]} has size {size}, which is not a whole multiple of the crop "
+                f"size {crop}. Call largest_multiple_center_crop() first, or use window_scan() "
+                f"which does this for you."
+            )
+
+    n_h, n_w, n_d = (s // c for s, c in zip(spatial_shape, (c_h, c_w, c_d)))
+
+    # Depth varies fastest, matching MONAI's GridPatch and the RoPE coordinate order.
+    crops = x.contiguous().view(channels, n_h, c_h, n_w, c_w, n_d, c_d)
+    crops = crops.permute(1, 3, 5, 0, 2, 4, 6)
+    crops = crops.reshape(n_h * n_w * n_d, channels, c_h, c_w, c_d)
+    return crops, (n_h, n_w, n_d)
+
+
+def _pad_short_axes(
+    x: torch.Tensor,
+    crop_size: Sequence[int],
+    pad_value: float,
+) -> torch.Tensor:
+    """Symmetrically pad any spatial axis shorter than its crop up to exactly one crop."""
+    spatial_shape = tuple(x.shape[1:])
+    short = [
+        (axis, size, int(crop))
+        for axis, (size, crop) in enumerate(zip(spatial_shape, crop_size))
+        if size < crop
+    ]
+    if not short:
+        return x
+
+    # F.pad takes the last dimension first, so build the spec back to front: D, W, H.
+    pad_spec: list = []
+    for axis in reversed(range(3)):
+        size, crop = spatial_shape[axis], int(crop_size[axis])
+        if size < crop:
+            total = crop - size
+            pad_spec.extend([total // 2, total - total // 2])
+        else:
+            pad_spec.extend([0, 0])
+
+    detail = ', '.join(
+        f"{'HWD'[axis]}={size} < {crop}" for axis, size, crop in short
+    )
+    warnings.warn(
+        f"Scan is smaller than one crop along {detail}. Padding with {pad_value} to reach the "
+        f"crop size {tuple(int(c) for c in crop_size)}. SPECTRE was not pretrained on padded "
+        f"volumes, so the resulting embedding is out of distribution. Pass pad_short_axes=False "
+        f"to raise instead.",
+        RuntimeWarning,
+        stacklevel=3,
+    )
+    return F.pad(x, pad_spec, mode='constant', value=pad_value)
+
+
+def window_scan(
+    x: torch.Tensor,
+    crop_size: Sequence[int] = DEFAULT_CROP_SIZE,
+    *,
+    scale_intensity: bool = True,
+    hu_range: Tuple[float, float] = DEFAULT_HU_RANGE,
+    pad_short_axes: bool = True,
+    pad_value: Optional[float] = None,
+) -> Tuple[torch.Tensor, Tuple[int, int, int]]:
+    """Turn a whole CT volume into the crops SPECTRE's backbone consumes.
+
+    Applies, in the order used for pretraining: HU scaling -> centre crop to a whole multiple of
+    `crop_size` -> grid patching.
+
+    Args:
+        x: A `(H, W, D)` or `(C, H, W, D)` volume. In raw Hounsfield Units unless
+            `scale_intensity=False`.
+        crop_size: Spatial size `(H, W, D)` of one crop.
+        scale_intensity: Map `hu_range` onto [0, 1] with clipping. Turn off if `x` is already
+            normalised.
+        hu_range: The `(a_min, a_max)` HU window. SPECTRE was pretrained with (-1000, 1000).
+        pad_short_axes: Pad axes shorter than one crop instead of raising.
+        pad_value: Value to pad with. Defaults to the scaled air value (0.0) when
+            `scale_intensity` is on, and to `x.min()` otherwise.
+
+    Returns:
+        The crops `(N, C, cH, cW, cD)` and the grid `(n_h, n_w, n_d)`.
+    """
+    if x.ndim == 3:
+        x = x.unsqueeze(0)
+    elif x.ndim != 4:
+        raise ValueError(
+            f"Expected a (H, W, D) or (C, H, W, D) volume, got shape {tuple(x.shape)}."
+        )
+
+    if len(crop_size) != 3:
+        raise ValueError(f"crop_size must have 3 elements (H, W, D), got {tuple(crop_size)}.")
+
+    if scale_intensity:
+        if float(x.min()) > -100.0:
+            warnings.warn(
+                f"Input does not look like Hounsfield Units (min={float(x.min()):.1f}). It will "
+                f"be rescaled from {hu_range} to [0, 1], which double-scales already-normalised "
+                f"data. Pass scale_intensity=False if the volume is already normalised.",
+                RuntimeWarning,
+                stacklevel=2,
+            )
+        x = scale_intensity_range(x, a_min=hu_range[0], a_max=hu_range[1])
+
+    if pad_value is None:
+        pad_value = 0.0 if scale_intensity else float(x.min())
+
+    too_short = [
+        f"{'HWD'[axis]}={int(size)} < {int(crop)}"
+        for axis, (size, crop) in enumerate(zip(x.shape[1:], crop_size))
+        if size < crop
+    ]
+    if too_short:
+        if not pad_short_axes:
+            raise ValueError(
+                f"Scan is smaller than one crop along {', '.join(too_short)}. Pass "
+                f"pad_short_axes=True to pad with air, or supply a scan covering at least "
+                f"{tuple(int(c) for c in crop_size)} voxels."
+            )
+        x = _pad_short_axes(x, crop_size, pad_value)
+
+    x = largest_multiple_center_crop(x, crop_size)
+    return grid_patch(x, crop_size)