Initial commit

README.md

@@ -0,0 +1,179 @@
+📢 [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.  
+
+📢 [2026-02-21] SPECTRE has been accepted for presentation at **CVPR 2026** (Denver, Colorado, USA)!  
+
+📢 [2026-01-20] [Semantic segmentation](https://github.com/cviviers/nnUNet) code and configurations using the nnUNet framework are now released!  
+
+
+# SPECTRE 👻👻👻
+
+<p align="center">
+  <a href="https://pypi.org/project/spectre-fm/"><img alt="PyPI Version" src="https://img.shields.io/pypi/v/spectre-fm?style=flat-square&label=version&cacheSeconds=0" /></a>
+  <a href="https://pypi.org/project/spectre-fm/"><img alt="Python Versions" src="https://img.shields.io/pypi/pyversions/spectre-fm?style=flat-square&cacheSeconds=0" /></a>
+  <a href="https://pypi.org/project/spectre-fm/"><img alt="Downloads per Month" src="https://img.shields.io/pypi/dm/spectre-fm?style=flat-square&label=downloads&cacheSeconds=0" /></a>
+  <a href="https://github.com/cclaess/SPECTRE/blob/main/LICENSE"><img alt="License" src="https://img.shields.io/github/license/cclaess/SPECTRE?style=flat-square&cacheSeconds=0" /></a>
+  <a href="https://huggingface.co/cclaess/SPECTRE"><img alt="Model weights" src="https://img.shields.io/badge/models-Hugging%20Face-yellow?style=flat-square&cacheSeconds=0" /></a>
+  <a href="https://arxiv.org/abs/2511.17209"><img alt="Paper" src="https://img.shields.io/badge/paper-arXiv-b31b1b?style=flat-square&cacheSeconds=0" /></a>
+</p>
+
+<p align="center">
+   <img src="imgs/method_overview.jpg" alt="SPECTRE architecture and pretraining strategies" width="600"/>
+</p>
+
+SPECTRE (**S**elf-Supervised & Cross-Modal **P**r**e**training for **CT** **R**epresentation **E**xtraction) is a **Transformer-based foundation model for 3D Computed Tomography (CT) scans**, trained using **self-supervised learning** (SSL) and **cross-modal vision–language alignment** (VLA). It provides rich and generalizable representations from medical imaging data, which can be fine-tuned for downstream tasks such as segmentation, classification, and anomaly detection.  
+
+SPECTRE has been trained on a large cohort of **open-source CT scans** of the **human abdomen and thorax**, as well as **paired radiology reports** and **Electronic Health Record data**, enabling it to capture representations that generalize across datasets and clinical settings.  
+
+This repository provides pretrained SPECTRE models together with tools for fine-tuning and evaluation.
+
+## 🧠 Pretrained Models
+The pretrained SPECTRE model can easily be imported as follows:
+
+```python
+from spectre import SpectreImageFeatureExtractor, MODEL_CONFIGS
+import torch
+
+config = MODEL_CONFIGS['spectre-large-pretrained']
+model = SpectreImageFeatureExtractor.from_config(config)
+model.eval()
+
+# 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
+
+patch_size = (128, 128, 64)
+pH, pW, pD = patch_size
+
+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)
+
+with torch.no_grad():
+    features = model(
+      x, 
+      grid_size=(
+        H // pH,
+        W // pW,
+        D // pD,
+      ),
+    )
+print("Features shape:", features.shape)
+```
+
+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)              |
+
+## 🩻 Segmentation (nnUNet)
+
+If you're looking for a nnUNet-based segmentation pipeline that uses SPECTRE as the backbone, see: https://github.com/cviviers/nnUNet
+
+## 📂 Repository Contents
+
+This repository is organized as follows:
+
+- 🚀 **`src/spectre/`** – Contains the core package, including:
+  - Pretraining methods
+  - Model architectures
+  - Data handling and transformations
+
+- 🛠️ **`src/spectre/configs/`** – Stores configuration files for different training settings.
+
+- 🔬 **`experiments/`** – Includes Python scripts for running various pretraining and downstream experiments.
+
+- 🐳 **`Dockerfile`** – Defines the environment for running a local version of SPECTRE inside a container.
+
+## ⚙️ Setting Up the Environment
+
+To get up and running with SPECTRE, install the base package with pip:
+
+```bash
+pip install spectre-fm
+```
+
+This installs only the runtime dependencies needed to load and run the pretrained models.
+
+If you want to fine-tune or pretrain SPECTRE, install the matching extra:
+
+```bash
+pip install "spectre-fm[training]"
+```
+
+If you only need the evaluation stack, install:
+
+```bash
+pip install "spectre-fm[eval]"
+```
+
+If training on GDS-enabled systems is required, install the CUDA 12 specific extra:
+
+```bash
+pip install "spectre-fm[gds-cuda12]"  # with training stack: "spectre-fm[training,gds-cuda12]"
+```
+
+**Note that** `gds-cuda12` is only compatible with CUDA 12.x environments.
+
+To install everything at once, use:
+
+```bash
+pip install "spectre-fm[all]"
+```
+
+or install the latest updates directly from GitHub:
+
+```bash
+pip install git+https://github.com/cclaess/SPECTRE.git
+```
+
+## 🐳 Building and Using Docker
+
+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
+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}},
+  url = {http://arxiv.org/abs/2511.17209},
+  doi = {10.48550/arXiv.2511.17209},
+  author = {Claessens, Cris and Viviers, Christiaan and D'Amicantonio, Giacomo and Bondarev, Egor and Sommen, Fons van der},
+  year={2025},
+}
+```
+
+## 🤝 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)

config.json

@@ -0,0 +1,32 @@
+{
+  "architectures": [
+    "SpectreModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_spectre.SpectreConfig",
+    "AutoModel": "modeling_spectre.SpectreModel"
+  },
+  "backbone_kwargs": {
+    "global_pool": "",
+    "init_values": 1.0,
+    "num_classes": 0,
+    "pos_embed": "rope",
+    "rope_kwargs": {
+      "base": 1000.0
+    }
+  },
+  "backbone_name": "vit_large_patch16_128",
+  "dtype": "float32",
+  "feature_combiner_kwargs": {
+    "global_pool": "",
+    "init_values": 1.0,
+    "num_classes": 0,
+    "pos_embed": "rope",
+    "rope_kwargs": {
+      "base": 100.0
+    }
+  },
+  "feature_combiner_name": "feat_vit_large",
+  "model_type": "spectre",
+  "transformers_version": "5.3.0"
+}

configuration_spectre.py

@@ -0,0 +1,32 @@
+from transformers import PretrainedConfig
+
+
+class SpectreConfig(PretrainedConfig):
+    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,
+        },
+        **kwargs,
+    ):
+        super().__init__(**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 {}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c90810e90d833cfff54880997fad4d964963fef4dd6824789f44a5bed063cb61
+size 1587720248

modeling_spectre.py

@@ -0,0 +1,38 @@
+import torch
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutput
+
+from spectre.model import SpectreImageFeatureExtractor
+from configuration_spectre import SpectreConfig
+
+
+class SpectreModel(PreTrainedModel):
+    config_class = SpectreConfig
+    base_model_prefix = "spectre"
+    main_input_name = "pixel_values"
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.model = SpectreImageFeatureExtractor(
+            backbone_name=config.backbone_name,
+            backbone_kwargs=config.backbone_kwargs,
+            feature_combiner_name=config.feature_combiner_name,
+            feature_combiner_kwargs=config.feature_combiner_kwargs,
+        )
+
+        self.post_init()
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        grid_size=None,
+        return_dict=True,
+        **kwargs,
+    ):
+        outputs = self.model(pixel_values, grid_size=grid_size)
+
+        if not return_dict:
+            return (outputs,)
+
+        return BaseModelOutput(last_hidden_state=outputs)

spectre/__init__.py

@@ -0,0 +1,29 @@
+"""Top-level package for spectre.
+
+Expose a small, stable public API here so users can do:
+
+	from spectre import SpectreImageFeatureExtractor, models
+
+Keep implementations in subpackages; this file only re-exports the most
+important symbols and subpackages for convenience.
+"""
+from .model import SpectreImageFeatureExtractor, MODEL_CONFIGS
+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",
+	"data",
+	"transforms",
+	"ssl",
+	"utils",
+	"__version__",
+    "__author__",
+    "__email__",
+]

spectre/model.py

@@ -0,0 +1,234 @@
+import torch
+import torch.nn as nn
+
+
+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.",
+    }
+}
+
+
+class SpectreImageFeatureExtractor(nn.Module):
+    def __init__(
+        self, 
+        backbone_name: str, 
+        backbone_kwargs: dict = {},
+        backbone_checkpoint_path_or_url: str | None = None,
+        feature_combiner_name: str | None = None, 
+        feature_combiner_kwargs: dict = {},
+        feature_combiner_checkpoint_path_or_url: str | None = None,
+        **kwargs,
+    ):  
+        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,
+            **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,
+                **kwargs,
+            )
+
+    def _init_backbone(
+        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, 
+            **kwargs,
+        )
+        
+    def _init_feature_combiner(
+        self, 
+        model_name: str,
+        checkpoint_path_or_url: str | None = None,
+        **kwargs,
+    ):
+        if self.backbone.global_pool == '':
+            patch_dim = self.backbone.embed_dim * 2  # CLS + AVG pooled tokens
+        else:
+            patch_dim = self.backbone.embed_dim
+
+        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, 
+            **kwargs,
+        )
+
+    def extract_backbone_features(
+        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 
+        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 
+        tokens and F is the feature dimension.
+
+        Args:
+            x (torch.Tensor): Input tensor of shape (B, N, C, H, W, D)
+        Returns:
+            torch.Tensor: Extracted features of shape (B, N, T, F)
+        """
+        assert x.ndim == 6, "Input tensor must have 6 dimensions: (B, N, C, H, W, D)"
+        B, N, C, H, W, D = x.shape
+        x = x.view(B * N, C, H, W, D)
+        features = self.backbone(x)
+        if features.ndim == 2:  # only CLS token
+            features = features.unsqueeze(1)
+        features = features.view(B, N, features.shape[1], -1)
+        return features
+
+    def combine_features(
+        self, 
+        features: torch.Tensor,
+        grid_size: tuple[int, int, int],
+    ):
+        """
+        Combine features from multiple image patches using the feature combiner.
+
+        Args:
+            features (torch.Tensor): Input features of shape (B, N, T, F)
+            grid_size (tuple[int, int, int]): Grid size of the image patches
+        Returns:
+            torch.Tensor: Combined features of shape (B, T', F')
+        """
+        _, N, T, _ = features.shape
+        assert features.ndim == 4, "Input features must have 4 dimensions: (B, N, T, F)"
+        assert N == grid_size[0] * grid_size[1] * grid_size[2], \
+            "Number of patches N must match the product of grid_size dimensions"
+
+        if T == 1:  # only CLS token
+            features = features.squeeze(2)
+        else:
+            # We combine CLS tokens with AVG pooling of other tokens
+            features = torch.cat([
+                features[:, :, 0, :],  # CLS token (B, N, F)
+                features[:, :, 1:, :].mean(dim=2)  # AVG pooled tokens (B, N, F)
+            ], dim=-1)  # (B, N, 2F)
+        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
+
+    @classmethod
+    def from_config(
+        cls, 
+        config: dict,
+        **kwargs,
+    ) -> 'SpectreImageFeatureExtractor':
+        
+        model = cls(
+            backbone_name=config["backbone"],
+            backbone_checkpoint_path_or_url=config.get("backbone_checkpoint_path_or_url", None),
+            backbone_kwargs=config.get("backbone_kwargs", {}),
+            feature_combiner_name=config.get("feature_combiner", None),
+            feature_combiner_checkpoint_path_or_url=config.get("feature_combiner_checkpoint_path_or_url", None),
+            feature_combiner_kwargs=config.get("feature_combiner_kwargs", {}),
+            **kwargs,
+        )
+        return model

spectre/models/__init__.py

@@ -0,0 +1,53 @@
+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,
+)
+from .resnet import (
+    ResNet,
+    resnet18,
+    resnet34, 
+    resnet50, 
+    resnet101, 
+    resnext50,
+    resnext101,
+)
+from .eomt import EoMT
+from .seomt import SEoMT
+from .upsample_anything import UPA
+
+__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',
+    'ResNet',
+    'resnet18',
+    'resnet34',
+    'resnet50',
+    'resnet101',
+    'resnext50',
+    'resnext101',
+    'EoMT',
+    'UPA',
+    'SEoMT',
+]

spectre/models/eomt.py

@@ -0,0 +1,230 @@
+# Adapted from https://github.com/tue-mps/eomt/
+from __future__ import annotations
+
+import math
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from spectre.models.layers import LayerNorm3d
+
+
+class ScaleBlock(nn.Module):
+    def __init__(
+        self, 
+        embed_dim: int, 
+        scale_factors: Union[int, Tuple[int, int, int]] = (2, 2, 2),
+        conv1_layer: nn.Module = nn.ConvTranspose3d,
+    ):
+        super().__init__()
+
+        self.conv1 = conv1_layer(
+            embed_dim,
+            embed_dim,
+            kernel_size=scale_factors,
+            stride=scale_factors,
+        )
+        self.act = nn.GELU()
+        self.conv2 = nn.Conv3d(
+            embed_dim,
+            embed_dim,
+            kernel_size=3,
+            padding=1,
+            groups=embed_dim,
+            bias=False,
+        )
+        self.norm = LayerNorm3d(embed_dim)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.act(x)
+        x = self.conv2(x)
+        x = self.norm(x)
+
+        return x
+
+
+def compute_upscale_stages(patch_size, min_size=4):
+    # Compute how many times to upscale per dimension
+    num_stages = []
+    for size in patch_size:
+        stages = max(0, int(math.log2(size)) - int(math.log2(min_size)))
+        num_stages.append(stages)
+    return num_stages
+
+
+class EoMT(nn.Module):
+    def __init__(
+        self,
+        backbone: "VisionTransformer",
+        num_classes: int,
+        num_q: int,
+        num_blocks=4,
+        masked_attn_enabled=True,
+    ):
+        super().__init__()
+        self.backbone = backbone
+        self.num_q = num_q
+        self.num_blocks = num_blocks
+        self.masked_attn_enabled = masked_attn_enabled
+
+        self.register_buffer("attn_mask_probs", torch.ones(num_blocks))
+
+        self.q = nn.Embedding(num_q, self.backbone.embed_dim)
+
+        self.class_head = nn.Linear(self.backbone.embed_dim, num_classes + 1)
+
+        self.mask_head = nn.Sequential(
+            nn.Linear(self.backbone.embed_dim, self.backbone.embed_dim),
+            nn.GELU(),
+            nn.Linear(self.backbone.embed_dim, self.backbone.embed_dim),
+            nn.GELU(),
+            nn.Linear(self.backbone.embed_dim, self.backbone.embed_dim),
+        )
+
+        patch_size = self.backbone.patch_embed.patch_size
+        num_upscale_stages = compute_upscale_stages(patch_size, min_size=4)
+
+        # Build per-stage scale factors list
+        max_stages = max(num_upscale_stages)
+        upscale_blocks = []
+        for stage_idx in range(max_stages):
+            # for each dimension, upscale by 2 only if this dimension still has
+            # remaining upscales at this stage
+            scale_factors = tuple(
+                2 if stage_idx < num_upscale_stages[dim] else 1
+                for dim in range(len(patch_size))
+            )
+            upscale_blocks.append(ScaleBlock(self.backbone.embed_dim,
+                                            scale_factors=scale_factors))
+
+        self.upscale = nn.Sequential(*upscale_blocks)
+
+    def _predict(self, x: torch.Tensor):
+        q = x[:, : self.num_q, :]
+
+        class_logits = self.class_head(q)
+
+        x = x[:, self.num_q + self.backbone.num_prefix_tokens :, :]
+        x = x.transpose(1, 2).reshape(
+            x.shape[0], -1, *self.backbone.patch_embed.grid_size
+        )
+
+        mask_logits = torch.einsum(
+            "bqc, bchwd -> bqhwd", self.mask_head(q), self.upscale(x)
+        )
+
+        return mask_logits, class_logits
+
+    @torch.compiler.disable
+    def _disable_attn_mask(self, attn_mask, prob):
+        if prob < 1:
+            random_queries = (
+                torch.rand(attn_mask.shape[0], self.num_q, device=attn_mask.device)
+                > prob
+            )
+            attn_mask[
+                :, : self.num_q, self.num_q + self.backbone.num_prefix_tokens :
+            ][random_queries] = True
+
+        return attn_mask
+
+    def _attn(self, module: 'Attention', x: torch.Tensor, mask: Optional[torch.Tensor], rope=None):
+        B, N, C = x.shape
+
+        q = module.q(x).reshape(B, N, module.num_heads, module.head_dim).permute(0, 2, 1, 3)
+        kv = module.kv(x).reshape(B, N, 2, module.num_heads, module.head_dim)
+        k, v = kv.permute(2, 0, 3, 1, 4).unbind(0)
+        q, k = module.q_norm(q), module.k_norm(k)
+
+        if mask is not None:
+            mask = mask[:, None, ...].expand(-1, module.num_heads, -1, -1)
+
+        dropout_p = module.attn_drop.p if self.training else 0.0
+
+        if rope is not None:
+            if isinstance(rope, list):
+                rope = tuple(torch.stack([r[i] for r in rope], dim=0) for i in range(2))
+            q, k = module.apply_rotary_pos_emb(q, k, rope)
+
+        if module.fused_attn:
+            x = F.scaled_dot_product_attention(q, k, v, mask, dropout_p)
+        else:
+            attn = (q @ k.transpose(-2, -1)) * module.scale
+            if mask is not None:
+                attn = attn.masked_fill(~mask, float("-inf"))
+            attn = F.softmax(attn, dim=-1)
+            attn = module.attn_drop(attn)
+            x = attn @ v
+
+        x = module.proj_drop(module.proj(x.transpose(1, 2).reshape(B, N, C)))
+
+        return x
+
+    def forward(self, x: torch.Tensor):
+        x = self.backbone.patch_embed(x)
+        x, rope = self.backbone._pos_embed(x)
+        x = self.backbone.patch_drop(x)
+        x = self.backbone.norm_pre(x)
+
+        attn_mask = None
+        mask_logits_per_layer, class_logits_per_layer = [], []
+
+        for i, block in enumerate(self.backbone.blocks):
+            if i == len(self.backbone.blocks) - self.num_blocks:
+                x = torch.cat(
+                    (self.q.weight[None, :, :].expand(x.shape[0], -1, -1), x), dim=1
+                )
+
+            if (
+                self.masked_attn_enabled
+                and i >= len(self.backbone.blocks) - self.num_blocks
+            ):
+                mask_logits, class_logits = self._predict(self.backbone.norm(x))
+                mask_logits_per_layer.append(mask_logits)
+                class_logits_per_layer.append(class_logits)
+
+                attn_mask = torch.ones(
+                    x.shape[0],
+                    x.shape[1],
+                    x.shape[1],
+                    dtype=torch.bool,
+                    device=x.device,
+                )
+                interpolated = F.interpolate(
+                    mask_logits,
+                    self.backbone.patch_embed.grid_size,
+                    mode="trilinear",
+                )
+                interpolated = interpolated.view(
+                    interpolated.size(0), interpolated.size(1), -1
+                )
+                attn_mask[
+                    :,
+                    : self.num_q,
+                    self.num_q + self.backbone.num_prefix_tokens :,
+                ] = (
+                    interpolated > 0
+                )
+                attn_mask = self._disable_attn_mask(
+                    attn_mask,
+                    self.attn_mask_probs[
+                        i - len(self.backbone.blocks) + self.num_blocks
+                    ],
+                )
+
+            x = x + block.drop_path1(
+                block.ls1(self._attn(block.attn, block.norm1(x), attn_mask, rope))
+            )
+            x = x + block.drop_path2(block.ls2(block.mlp(block.norm2(x))))
+
+        mask_logits, class_logits = self._predict(self.backbone.norm(x))
+        mask_logits_per_layer.append(mask_logits)
+        class_logits_per_layer.append(class_logits)
+
+        return (
+            mask_logits_per_layer,
+            class_logits_per_layer,
+        )

spectre/models/layers/__init__.py

@@ -0,0 +1,11 @@
+from .patch_embed import PatchEmbed
+from .attention import Attention
+from .layernorm import LayerNorm3d
+from .rotary_pos_embed import RotaryPositionEmbedding
+
+__all__ = [
+    'PatchEmbed',
+    'Attention',
+    'LayerNorm3d',
+    'RotaryPositionEmbedding',
+]

spectre/models/layers/attention.py

@@ -0,0 +1,187 @@
+from typing import Type, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.jit import Final
+from timm.layers import use_fused_attn
+
+from spectre.models.layers.rotary_pos_embed import rope_apply
+
+
+class Attention(nn.Module):
+    fused_attn: Final[bool]
+
+    def __init__(
+            self,
+            dim: int,
+            num_heads: int = 8,
+            mode: str = "mha",
+            q_proj_dim: Optional[int] = None,
+            kv_proj_dim: Optional[int] = None,
+            qkv_bias: bool = False,
+            qk_norm: bool = False,
+            proj_bias: bool = True,
+            attn_drop: float = 0.,
+            proj_drop: float = 0.,
+            norm_layer: Type[nn.Module] = nn.LayerNorm,
+    ) -> None:
+        super().__init__()
+        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim ** -0.5
+        self.fused_attn = use_fused_attn()
+        self.mode = mode.lower()
+        assert self.mode in ["mha", "mqa", "mla"], "Attention mode must be 'mha', 'mqa', or 'mla'"
+        assert not (self.mode == "mla" and kv_proj_dim is None), "kv_proj_dim must be provided for 'mla' mode"
+        assert not (self.mode == "mla" and q_proj_dim is None), "q_proj_dim must be provided for 'mla' mode"
+        
+        if self.mode == "mha":
+            self.q = nn.Linear(dim, dim, bias=qkv_bias)
+            self.kv = nn.Linear(dim, 2 * dim, bias=qkv_bias)  # Key and value pair for every head
+            self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+            self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        elif self.mode == "mqa":
+            self.q = nn.Linear(dim, dim, bias=qkv_bias)
+            self.kv = nn.Linear(dim, 2 * self.head_dim, bias=qkv_bias)  # Key and value pair shared across heads
+            self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+            self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        elif self.mode == "mla":
+            self.q_proj = nn.Linear(dim, q_proj_dim, bias=qkv_bias)  # Projected query for every head
+            self.kv_proj = nn.Linear(dim, kv_proj_dim, bias=qkv_bias)  # Projected key and value pair for every head
+            self.q_norm = norm_layer(q_proj_dim) if qk_norm else nn.Identity()
+            self.kv_norm = norm_layer(kv_proj_dim) if qk_norm else nn.Identity()
+            self.q = nn.Linear(q_proj_dim, dim, bias=qkv_bias)  # Query for every head
+            self.kv = nn.Linear(kv_proj_dim, 2 * dim, bias=qkv_bias)  # Key and value pair for every head
+        
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim, bias=proj_bias)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def apply_rotary_pos_emb(
+        self, 
+        q: torch.Tensor, 
+        k: torch.Tensor, 
+        rope: Tuple[torch.Tensor, torch.Tensor],
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Apply RoPE to the query and key tensors.
+        Args:
+            q (torch.Tensor): Query tensor of shape (B, num_heads, N, head_dim)
+            k (torch.Tensor): Key tensor of shape (B, num_heads, N, head_dim)
+            rope (Tuple[torch.Tensor, torch.Tensor]): Tuple of (sin, cos) tensors for RoPE application.
+                Sin and cos can be of shape 
+        """
+        
+        # Match dtype to rope for numeric stability
+        q_dtype, k_dtype = q.dtype, k.dtype
+        sin, cos = rope
+
+        if sin.ndim == 2:
+            sin = sin.unsqueeze(0).unsqueeze(0)
+            cos = cos.unsqueeze(0).unsqueeze(0)
+        elif sin.ndim == 3:
+            sin = sin.unsqueeze(1)
+            cos = cos.unsqueeze(1)
+        else:
+            raise ValueError("RoPE sin/cos must be of shape [N, head_dim] or [B, N, head_dim]")
+        
+        rope_dtype = sin.dtype
+
+        q = q.to(dtype=rope_dtype)
+        k = k.to(dtype=rope_dtype)
+
+        N = q.shape[-2]             # total tokens per sample
+        N_spatial = sin.shape[-2]   # number of spatial tokens covered by rope
+        prefix = N - N_spatial      # e.g., [cls] or [reg] tokens at the front
+        assert prefix >= 0, "RoPE sin/cos length exceeds sequence length"
+
+        if prefix > 0:
+            q_prefix = q[:, :, :prefix, :]
+            k_prefix = k[:, :, :prefix, :]
+            q_spatial = q[:, :, prefix:, :]
+            k_spatial = k[:, :, prefix:, :]
+        else:
+            q_prefix = k_prefix = None
+            q_spatial, k_spatial = q, k
+
+        # Apply RoPE on the spatial tail
+        q_spatial = rope_apply(q_spatial, sin, cos)
+        k_spatial = rope_apply(k_spatial, sin, cos)
+
+        # Stitch back
+        if prefix > 0:
+            q = torch.cat((q_prefix, q_spatial), dim=-2)
+            k = torch.cat((k_prefix, k_spatial), dim=-2)
+        else:
+            q, k = q_spatial, k_spatial
+
+        # Cast back to original dtypes
+        q = q.to(dtype=q_dtype)
+        k = k.to(dtype=k_dtype)
+
+        return q, k
+
+    def compute_qkv(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        B, N, _ = x.shape
+        if self.mode == "mha":
+            q = self.q(x).reshape(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
+            kv = self.kv(x).reshape(B, N, 2, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+            k, v = kv.unbind(0)
+            q, k = self.q_norm(q), self.k_norm(k)
+        elif self.mode == "mqa":
+            q = self.q(x).reshape(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
+            kv = self.kv(x).reshape(B, N, 2, 1, self.head_dim).permute(2, 0, 3, 1, 4)
+            kv = kv.expand(-1, -1, self.num_heads, -1, -1)  # Expand to match num_heads
+            k, v = kv.unbind(0)
+            q, k = self.q_norm(q), self.k_norm(k)
+        elif self.mode == "mla":
+            q = self.q_proj(x)
+            kv = self.kv_proj(x)
+            q, kv = self.q_norm(q), self.kv_norm(kv)  # Normalization on projections
+            q = self.q(q).reshape(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
+            kv = self.kv(kv).reshape(B, N, 2, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+            k, v = kv.unbind(0)
+        return q, k, v
+        
+    def compute_attention(
+        self, 
+        q: torch.Tensor, 
+        k: torch.Tensor, 
+        v: torch.Tensor,
+    ) -> torch.Tensor:
+        B, _, N, _ = q.shape
+        C = self.num_heads * self.head_dim
+
+        if self.fused_attn:
+            x = F.scaled_dot_product_attention(
+                q, k, v,
+                dropout_p=self.attn_drop.p if self.training else 0.,
+            )
+        else:
+            q = q * self.scale
+            attn = q @ k.transpose(-2, -1)
+            attn = attn.softmax(dim=-1)
+            attn = self.attn_drop(attn)
+            x = attn @ v
+
+        return x.transpose(1, 2).reshape(B, N, C)
+
+    def forward(
+        self, 
+        x: torch.Tensor, 
+        rope = None,
+    ) -> torch.Tensor:
+
+        q, k, v = self.compute_qkv(x)
+
+        if rope is not None:
+            if isinstance(rope, list):
+                rope = tuple(torch.stack([r[i] for r in rope], dim=0) for i in range(2))
+            q, k = self.apply_rotary_pos_emb(q, k, rope)
+        
+        x = self.compute_attention(q, k, v)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x

spectre/models/layers/layernorm.py

@@ -0,0 +1,23 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from timm.layers.fast_norm import is_fast_norm, fast_layer_norm
+
+
+class LayerNorm3d(nn.LayerNorm):
+    """ LayerNorm for channels of '3D' spatial NCHWD tensors """
+    _fast_norm: torch.jit.Final[bool]
+
+    def __init__(self, num_channels, eps=1e-6, affine=True):
+        super().__init__(num_channels, eps=eps, elementwise_affine=affine)
+        self._fast_norm = is_fast_norm()  # Assuming is_fast_norm() is defined somewhere
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x.permute(0, 2, 3, 4, 1)  # Permute to NCDHW format
+        if self._fast_norm:
+            x = fast_layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        else:
+            x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        x = x.permute(0, 4, 1, 2, 3)  # Permute back to NCHWD format
+        return x
+    

spectre/models/layers/patch_embed.py

@@ -0,0 +1,135 @@
+import math
+from typing import Optional, Tuple, Union, Callable
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from spectre.utils import (
+    to_3tuple,
+    resample_patch_embed, 
+    Format, 
+    nchwd_to,
+)
+
+
+class PatchEmbed(nn.Module):
+    """ 3D Image to Patch Embedding
+    """
+    output_fmt: Format
+    dynamic_img_pad: torch.jit.Final[bool]
+
+    def __init__(
+        self,
+        img_size: Optional[Union[int, Tuple[int, int, int]]] = (128, 128, 64),
+        patch_size: Union[int, Tuple[int, int, int]] = (16, 16, 8),
+        in_chans: int = 1,
+        embed_dim: int = 768,
+        norm_layer: Optional[Callable] = None,
+        flatten: bool = True,
+        output_fmt: Optional[str] = None,
+        bias: bool = True,
+        strict_img_size: bool = True,
+        dynamic_img_pad: bool = False,
+    ):
+        super().__init__()
+        self.patch_size = to_3tuple(patch_size)
+        self.img_size, self.grid_size, self.num_patches = self._init_img_size(img_size)
+
+        if output_fmt is not None:
+            self.flatten = False
+            self.output_fmt = Format(output_fmt)
+        else:
+            # flatten spatial dim and transpose to channels last, kept for bwd compat
+            self.flatten = flatten
+            self.output_fmt = Format.NCHWD
+        self.strict_img_size = strict_img_size
+        self.dynamic_img_pad = dynamic_img_pad
+
+        self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+    
+    def _init_img_size(self, img_size: Union[int, Tuple[int, int, int]]):
+        assert self.patch_size
+        if img_size is None:
+            return None, None, None
+        img_size = to_3tuple(img_size)
+        grid_size = tuple([s // p for s, p in zip(img_size, self.patch_size)])
+        num_patches = grid_size[0] * grid_size[1] * grid_size[2]
+        return img_size, grid_size, num_patches
+    
+    def set_input_size(
+            self,
+        img_size: Optional[Union[int, Tuple[int, int, int]]] = None,
+        patch_size: Optional[Union[int, Tuple[int, int, int]]] = None,
+    ):
+        new_patch_size = None
+        if patch_size is not None:
+            new_patch_size = to_3tuple(patch_size)
+        if new_patch_size is not None and new_patch_size != self.patch_size:
+            with torch.no_grad():
+                new_proj = nn.Conv3d(
+                    self.proj.in_channels,
+                    self.proj.out_channels,
+                    kernel_size=new_patch_size,
+                    stride=new_patch_size,
+                    bias=self.proj.bias is not None,
+                )
+                new_proj.weight.copy_(resample_patch_embed(self.proj.weight, new_patch_size, verbose=True))
+                if self.proj.bias is not None:
+                    new_proj.bias.copy_(self.proj.bias)
+                self.proj = new_proj
+            self.patch_size = new_patch_size
+        img_size = img_size or self.img_size
+        if img_size != self.img_size or new_patch_size is not None:
+            self.img_size, self.grid_size, self.num_patches = self._init_img_size(img_size)
+    
+    def feat_ratio(self, as_scalar=True) -> Union[Tuple[int, int, int], int]:
+        if as_scalar:
+            return max(self.patch_size)
+        else:
+            return self.patch_size
+        
+    def dynamic_feat_size(self, img_size: Tuple[int, int, int]) -> Tuple[int, int, int]:
+        """ Get grid (feature) size for given image size taking account of dynamic padding.
+        NOTE: must be torchscript compatible so using fixed tuple indexing
+        """
+        if self.dynamic_img_pad:
+            return (
+                math.ceil(img_size[0] / self.patch_size[0]), 
+                math.ceil(img_size[1] / self.patch_size[1]),
+                math.ceil(img_size[2] / self.patch_size[2]),
+            )
+        else:
+            return (
+                img_size[0] // self.patch_size[0], 
+                img_size[1] // self.patch_size[1],
+                img_size[2] // self.patch_size[2],
+            )
+
+    def forward(self, x):
+        _, _, H, W, D = x.shape
+        if self.img_size is not None:
+            if self.strict_img_size:
+                assert H == self.img_size[0], f"Input height ({H}) doesn't match model ({self.img_size[0]})."
+                assert W == self.img_size[1], f"Input width ({W}) doesn't match model ({self.img_size[1]})."
+                assert D == self.img_size[2], f"Input depth ({D}) doesn't match model ({self.img_size[2]})."
+            elif not self.dynamic_img_pad:
+                assert H % self.patch_size[0] == 0, \
+                    f"Input height ({H}) should be divisible by patch size ({self.patch_size[0]})."
+                assert W % self.patch_size[1] == 0, \
+                    f"Input width ({W}) should be divisible by patch size ({self.patch_size[1]})."
+                assert D % self.patch_size[2] == 0, \
+                    f"Input depth ({D}) should be divisible by patch size ({self.patch_size[2]})."
+        if self.dynamic_img_pad:
+            pad_h = (self.patch_size[0] - H % self.patch_size[0]) % self.patch_size[0]
+            pad_w = (self.patch_size[1] - W % self.patch_size[1]) % self.patch_size[1]
+            pad_d = (self.patch_size[2] - D % self.patch_size[2]) % self.patch_size[2]
+            x = F.pad(x, (0, pad_w, 0, pad_h, 0, pad_d))
+        x = self.proj(x)
+        if self.flatten:
+            x = x.flatten(2).transpose(1, 2)  # NCHWD -> NLC
+        elif self.output_fmt != Format.NCHWD:
+            x = nchwd_to(x, self.output_fmt)
+        x = self.norm(x)
+        return x

spectre/models/layers/rotary_pos_embed.py

@@ -0,0 +1,157 @@
+import math
+from typing import Literal, Tuple
+
+import torch
+import torch.nn as nn
+import numpy as np
+
+
+def rope_rotate_half(x: torch.Tensor) -> torch.Tensor:
+    # x:   [..., D], split into halves and rotate
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat([-x2, x1], dim=-1)
+
+
+def rope_apply(
+    x: torch.Tensor, 
+    sin: torch.Tensor, 
+    cos: torch.Tensor
+) -> torch.Tensor:
+    # x, sin, cos: [..., D]
+    return (x * cos) + (rope_rotate_half(x) * sin)
+
+
+class RotaryPositionEmbedding(nn.Module):
+    """
+    3D Rotary Positional Embedding (RoPE) with no mixing across axes (axial),
+    and no learnable weights. Allows for shifting and scaling of the positional encodings
+    for improving performance on varying resolutions.
+    Mirrors DINOv3 style but for (H, W, D).
+
+    Requirements:
+      - head_dim % 6 == 0  (because 3 axes -> periods of size head_dim//6, then we tile to fill head_dim)
+
+    Two parametrizations:
+      * base
+      * min_period + max_period
+    """
+    def __init__(
+        self,
+        embed_dim: int,
+        *,
+        num_heads: int,
+        base: float | None = 1000.0,  # works for common 8^3=512 to 16^3=4096 tokens
+        min_period: float | None = None,
+        max_period: float | None = None,
+        normalize_coords: Literal["min", "max", "separate"] = "separate",
+        shift_coords: float | None = None,
+        jitter_coords: float | None = None,
+        rescale_coords: float | None = None,
+        dtype: torch.dtype | None = None,
+        device: torch.device | None = None,
+    ):
+        super().__init__()
+        assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
+        head_dim = embed_dim // num_heads
+        assert head_dim % 6 == 0, "For 3D RoPE, (embed_dim // num_heads) must be divisible by 6"
+
+        both_periods = (min_period is not None) and (max_period is not None)
+        if (base is None and not both_periods) or (base is not None and both_periods):
+            raise ValueError("Either `base` or `min_period`+`max_period` must be provided.")
+
+        self.base = base
+        self.min_period = min_period
+        self.max_period = max_period
+        self.head_dim = head_dim
+        self.normalize_coords = normalize_coords
+        self.shift_coords = shift_coords
+        self.jitter_coords = jitter_coords
+        self.rescale_coords = rescale_coords
+
+        # Keep dtype persistent so teacher can be initialized from student state_dict()
+        self.dtype = dtype
+        self.register_buffer(
+            "periods",
+            torch.empty(self.head_dim // 6, device=device, dtype=dtype),
+            persistent=True,
+        )
+        self._init_weights()
+
+    @torch.no_grad()
+    def _init_weights(self):
+        device = self.periods.device
+        dtype = self.dtype
+        if self.base is not None:
+            # powers from 0..(head_dim // 3 - 1), normalized to [0, 1) across head_dim // 3? 
+            # for 3D we use // 6 per axis
+            periods = self.base ** (
+                2 * torch.arange(self.head_dim // 6, device=device, dtype=dtype) / (self.head_dim // 3)
+            )
+        else:
+            # geometric spacing between min_period and max_period
+            base = self.max_period / self.min_period
+            exponents = torch.linspace(0, 1, self.head_dim // 6, device=device, dtype=dtype)
+            periods = base ** exponents
+            periods = periods / base
+            periods = periods * self.max_period
+        self.periods.data = periods
+
+    def forward(self, *, H: int, W: int, D: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Returns:
+          sin, cos: [H * W * D, head_dim]   (per-head)
+        """
+        device = self.periods.device
+        dtype = self.dtype
+        dd = dict(device=device, dtype=dtype)
+
+        # Prepare coords in [0, 1] then map to [-1, +1]
+        if self.normalize_coords == "max":
+            max_dim = max(H, W, D)
+            coords_h = torch.arange(0.5, H, **dd) / max_dim
+            coords_w = torch.arange(0.5, W, **dd) / max_dim
+            coords_d = torch.arange(0.5, D, **dd) / max_dim
+        elif self.normalize_coords == "min":
+            min_dim = min(H, W, D)
+            coords_h = torch.arange(0.5, H, **dd) / min_dim
+            coords_w = torch.arange(0.5, W, **dd) / min_dim
+            coords_d = torch.arange(0.5, D, **dd) / min_dim
+        elif self.normalize_coords == "separate":
+            coords_h = torch.arange(0.5, H, **dd) / H
+            coords_w = torch.arange(0.5, W, **dd) / W
+            coords_d = torch.arange(0.5, D, **dd) / D
+        else:
+            raise ValueError(f"Unknown normalize_coords: {self.normalize_coords}")
+
+        coords = torch.stack(
+            torch.meshgrid(coords_h, coords_w, coords_d, indexing="ij"),
+            dim=-1
+        )  # [H, W, D, 3]
+        coords = coords.flatten(0, 2)                  # [HWD, 3]
+        coords = 2.0 * coords - 1.0                    # [-1, +1]
+
+        # Optional train-time augmentations on coords (DINOv3)
+        if self.training and self.shift_coords is not None:
+            shift_hwd = torch.empty(3, **dd).uniform_(-self.shift_coords, self.shift_coords)
+            coords = coords + shift_hwd[None, :]
+
+        if self.training and self.jitter_coords is not None:
+            jit_max = np.log(self.jitter_coords); jit_min = -jit_max
+            jitter = torch.empty(3, **dd).uniform_(jit_min, jit_max).exp()
+            coords = coords * jitter[None, :]
+
+        if self.training and self.rescale_coords is not None:
+            r_max = np.log(self.rescale_coords); r_min = -r_max
+            rescale = torch.empty(1, **dd).uniform_(r_min, r_max).exp()
+            coords = coords * rescale
+
+        # --- Build angles per axis, then concatenate across axes ---
+        # coords: [N, 3] ; periods: [head_dim // 6]
+        # angles: [N, 3, head_dim // 6] -> flatten(1, 2) -> [N, head_dim // 2] -> tile(2) -> [N, head_dim]
+        angles = 2 * math.pi * coords[:, :, None] / self.periods[None, None, :]  # [N, 3, head_dim // 6]
+        angles = angles.flatten(1, 2)                                            # [N, head_dim // 2]
+        angles = angles.tile(2)                                                  # [N, head_dim]
+
+        cos = torch.cos(angles)
+        sin = torch.sin(angles)
+        return sin, cos

spectre/models/resnet.py

@@ -0,0 +1,726 @@
+import os
+import math
+from urllib.parse import urlparse
+from typing import Type, Any, Tuple, List, Optional, Union, Dict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from spectre.utils import to_ntuple
+
+
+def get_padding(kernel_size: int, stride: int, dilation: int = 1) -> int:
+    padding = ((stride - 1) + dilation * (kernel_size - 1)) // 2
+    return padding
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(
+            self,
+            inplanes: int,
+            planes: int,
+            stride: int = 1,
+            downsample: Optional[nn.Module] = None,
+            cardinality: int = 1,
+            base_width: int = 64,
+            reduce_first: int = 1,
+            dilation: int = 1,
+            first_dilation: Optional[int] = None,
+            act_layer: Type[nn.Module] = nn.ReLU,
+            norm_layer: Type[nn.Module] = nn.BatchNorm3d,
+    ):
+        """
+        Args:
+            inplanes: Input channel dimensionality.
+            planes: Used to determine output channel dimensionalities.
+            stride: Stride used in convolution layers.
+            downsample: Optional downsample layer for residual path.
+            cardinality: Number of convolution groups.
+            base_width: Base width used to determine output channel dimensionality.
+            reduce_first: Reduction factor for first convolution output width of residual blocks.
+            dilation: Dilation rate for convolution layers.
+            first_dilation: Dilation rate for first convolution layer.
+            act_layer: Activation layer.
+            norm_layer: Normalization layer.
+        """
+        super(BasicBlock, self).__init__()
+
+        assert cardinality == 1, 'BasicBlock only supports cardinality of 1'
+        assert base_width == 64, 'BasicBlock does not support changing base width'
+        first_planes = planes // reduce_first
+        outplanes = planes * self.expansion
+        first_dilation = first_dilation or dilation
+
+        self.conv1 = nn.Conv3d(
+            inplanes, first_planes, kernel_size=3, stride=stride, padding=first_dilation,
+            dilation=first_dilation, bias=False)
+        self.bn1 = norm_layer(first_planes)
+        self.act1 = act_layer()
+
+        self.conv2 = nn.Conv3d(
+            first_planes, outplanes, kernel_size=3, padding=dilation, dilation=dilation, bias=False)
+        self.bn2 = norm_layer(outplanes)
+
+        self.act2 = act_layer()
+        self.downsample = downsample
+        self.stride = stride
+        self.dilation = dilation
+
+    def zero_init_last(self):
+        if getattr(self.bn2, 'weight', None) is not None:
+            nn.init.zeros_(self.bn2.weight)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        shortcut = x
+
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.act1(x)
+
+        x = self.conv2(x)
+        x = self.bn2(x)
+
+        if self.downsample is not None:
+            shortcut = self.downsample(shortcut)
+        x = x + shortcut
+        x = self.act2(x)
+
+        return x
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(
+            self,
+            inplanes: int,
+            planes: int,
+            stride: int = 1,
+            downsample: Optional[nn.Module] = None,
+            cardinality: int = 1,
+            base_width: int = 64,
+            reduce_first: int = 1,
+            dilation: int = 1,
+            first_dilation: Optional[int] = None,
+            act_layer: Type[nn.Module] = nn.ReLU,
+            norm_layer: Type[nn.Module] = nn.BatchNorm3d,
+    ):
+        """
+        Args:
+            inplanes: Input channel dimensionality.
+            planes: Used to determine output channel dimensionalities.
+            stride: Stride used in convolution layers.
+            downsample: Optional downsample layer for residual path.
+            cardinality: Number of convolution groups.
+            base_width: Base width used to determine output channel dimensionality.
+            reduce_first: Reduction factor for first convolution output width of residual blocks.
+            dilation: Dilation rate for convolution layers.
+            first_dilation: Dilation rate for first convolution layer.
+            act_layer: Activation layer.
+            norm_layer: Normalization layer.
+        """
+        super(Bottleneck, self).__init__()
+
+        width = int(math.floor(planes * (base_width / 64)) * cardinality)
+        first_planes = width // reduce_first
+        outplanes = planes * self.expansion
+        first_dilation = first_dilation or dilation
+
+        self.conv1 = nn.Conv3d(inplanes, first_planes, kernel_size=1, bias=False)
+        self.bn1 = norm_layer(first_planes)
+        self.act1 = act_layer()
+
+        self.conv2 = nn.Conv3d(
+            first_planes, width, kernel_size=3, stride=stride,
+            padding=first_dilation, dilation=first_dilation, groups=cardinality, bias=False)
+        self.bn2 = norm_layer(width)
+        self.act2 = act_layer()
+
+        self.conv3 = nn.Conv3d(width, outplanes, kernel_size=1, bias=False)
+        self.bn3 = norm_layer(outplanes)
+
+        self.act3 = act_layer()
+        self.downsample = downsample
+        self.stride = stride
+        self.dilation = dilation
+
+    def zero_init_last(self):
+        if getattr(self.bn3, 'weight', None) is not None:
+            nn.init.zeros_(self.bn3.weight)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        shortcut = x
+
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.act1(x)
+
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.act2(x)
+
+        x = self.conv3(x)
+        x = self.bn3(x)
+
+        if self.downsample is not None:
+            shortcut = self.downsample(shortcut)
+        x = x + shortcut
+        x = self.act3(x)
+
+        return x
+
+
+def downsample_conv(
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        dilation: int = 1,
+        first_dilation: Optional[int] = None,
+        norm_layer: Optional[Type[nn.Module]] = None,
+) -> nn.Module:
+    norm_layer = norm_layer or nn.BatchNorm3d
+    kernel_size = 1 if stride == 1 and dilation == 1 else kernel_size
+    first_dilation = (first_dilation or dilation) if kernel_size > 1 else 1
+    p = get_padding(kernel_size, stride, first_dilation)
+
+    return nn.Sequential(*[
+        nn.Conv3d(
+            in_channels, out_channels, kernel_size, stride=stride, padding=p, dilation=first_dilation, bias=False),
+        norm_layer(out_channels)
+    ])
+
+
+def downsample_avg(
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        dilation: int = 1,
+        first_dilation: Optional[int] = None,
+        norm_layer: Optional[Type[nn.Module]] = None,
+) -> nn.Module:
+    norm_layer = norm_layer or nn.BatchNorm3d
+    avg_stride = stride if dilation == 1 else 1
+    if stride == 1 and dilation == 1:
+        pool = nn.Identity()
+    else:
+        pool = nn.AvgPool3d(2, avg_stride, ceil_mode=True, count_include_pad=False)
+
+    return nn.Sequential(*[
+        pool,
+        nn.Conv3d(in_channels, out_channels, 1, stride=1, padding=0, bias=False),
+        norm_layer(out_channels)
+    ])
+
+
+def make_blocks(
+        block_fns: Tuple[Union[BasicBlock, Bottleneck]],
+        channels: Tuple[int, ...],
+        block_repeats: Tuple[int, ...],
+        inplanes: int,
+        reduce_first: int = 1,
+        output_stride: int = 32,
+        down_kernel_size: int = 1,
+        avg_down: bool = False,
+        **kwargs,
+) -> Tuple[List[Tuple[str, nn.Module]], List[Dict[str, Any]]]:
+    stages = []
+    feature_info = []
+    net_num_blocks = sum(block_repeats)
+    net_block_idx = 0
+    net_stride = 4
+    dilation = prev_dilation = 1
+    for stage_idx, (block_fn, planes, num_blocks) in enumerate(zip(block_fns, channels, block_repeats)):
+        stage_name = f'layer{stage_idx + 1}'  # never liked this name, but weight compat requires it
+        stride = 1 if stage_idx == 0 else 2
+        if net_stride >= output_stride:
+            dilation *= stride
+            stride = 1
+        else:
+            net_stride *= stride
+
+        downsample = None
+        if stride != 1 or inplanes != planes * block_fn.expansion:
+            down_kwargs = dict(
+                in_channels=inplanes,
+                out_channels=planes * block_fn.expansion,
+                kernel_size=down_kernel_size,
+                stride=stride,
+                dilation=dilation,
+                first_dilation=prev_dilation,
+                norm_layer=kwargs.get('norm_layer'),
+            )
+            downsample = downsample_avg(**down_kwargs) if avg_down else downsample_conv(**down_kwargs)
+
+        block_kwargs = dict(reduce_first=reduce_first, dilation=dilation, **kwargs)
+        blocks = []
+        for block_idx in range(num_blocks):
+            downsample = downsample if block_idx == 0 else None
+            stride = stride if block_idx == 0 else 1
+            blocks.append(block_fn(
+                inplanes,
+                planes,
+                stride,
+                downsample,
+                first_dilation=prev_dilation,
+                **block_kwargs,
+            ))
+            prev_dilation = dilation
+            inplanes = planes * block_fn.expansion
+            net_block_idx += 1
+
+        stages.append((stage_name, nn.Sequential(*blocks)))
+        feature_info.append(dict(num_chs=inplanes, reduction=net_stride, module=stage_name))
+
+    return stages, feature_info
+
+
+def feature_take_indices(
+        num_features: int,
+        indices: Optional[Union[int, List[int]]] = None,
+        as_set: bool = False,
+) -> Tuple[List[int], int]:
+    """ Determine the absolute feature indices to 'take' from.
+
+    Note: This function can be called in forward() so must be torchscript compatible,
+    which requires some incomplete typing and workaround hacks.
+
+    Args:
+        num_features: total number of features to select from
+        indices: indices to select,
+          None -> select all
+          int -> select last n
+          list/tuple of int -> return specified (-ve indices specify from end)
+        as_set: return as a set
+
+    Returns:
+        List (or set) of absolute (from beginning) indices, Maximum index
+    """
+    if indices is None:
+        indices = num_features  # all features if None
+
+    if isinstance(indices, int):
+        # convert int -> last n indices
+        assert 0 < indices <= num_features, f'last-n ({indices}) is out of range (1 to {num_features})'
+        take_indices = [num_features - indices + i for i in range(indices)]
+    else:
+        take_indices: List[int] = []
+        for i in indices:
+            idx = num_features + i if i < 0 else i
+            assert 0 <= idx < num_features, f'feature index {idx} is out of range (0 to {num_features - 1})'
+            take_indices.append(idx)
+
+    if not torch.jit.is_scripting() and as_set:
+        return set(take_indices), max(take_indices)
+
+    return take_indices, max(take_indices)
+
+
+class ResNet(nn.Module):
+    """ResNet / ResNeXt
+
+    This class implements all variants of ResNet, ResNeXt that
+      * have > 1 stride in the 3x3 conv layer of bottleneck
+      * have conv-bn-act ordering
+
+    This ResNet impl supports a number of stem and downsample options based on the v1c, v1d, v1e, and v1s
+    variants included in the MXNet Gluon ResNetV1b model. The C and D variants are also discussed in the
+    'Bag of Tricks' paper: https://arxiv.org/pdf/1812.01187. The B variant is equivalent to torchvision default.
+
+    ResNet variants (the same modifications can be used in SE/ResNeXt models as well):
+      * normal, b - 7x7 stem, stem_width = 64, same as torchvision ResNet, NVIDIA ResNet 'v1.5', Gluon v1b
+      * c - 3 layer deep 3x3 stem, stem_width = 32 (32, 32, 64)
+      * d - 3 layer deep 3x3 stem, stem_width = 32 (32, 32, 64), average pool in downsample
+      * e - 3 layer deep 3x3 stem, stem_width = 64 (64, 64, 128), average pool in downsample
+      * s - 3 layer deep 3x3 stem, stem_width = 64 (64, 64, 128)
+      * t - 3 layer deep 3x3 stem, stem width = 32 (24, 48, 64), average pool in downsample
+      * tn - 3 layer deep 3x3 stem, stem width = 32 (24, 32, 64), average pool in downsample
+
+    ResNeXt
+      * normal - 7x7 stem, stem_width = 64, standard cardinality and base widths
+      * same c,d, e, s variants as ResNet can be enabled
+    """
+
+    def __init__(
+            self,
+            block: Union[BasicBlock, Bottleneck],
+            layers: Tuple[int, ...],
+            num_classes: int = 1000,
+            in_chans: int = 1,
+            output_stride: int = 32,
+            global_pool: str = 'avg',
+            cardinality: int = 1,
+            base_width: int = 64,
+            stem_width: int = 64,
+            stem_type: str = '',
+            replace_stem_pool: bool = False,
+            block_reduce_first: int = 1,
+            down_kernel_size: int = 1,
+            avg_down: bool = False,
+            channels: Optional[Tuple[int, ...]] = (64, 128, 256, 512),
+            act_layer: Type[nn.Module] = nn.ReLU,
+            norm_layer: Type[nn.Module] = nn.BatchNorm3d,
+            drop_rate: float = 0.0,
+            zero_init_last: bool = True,
+            block_args: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        Args:
+            block (nn.Module): class for the residual block. Options are BasicBlock, Bottleneck.
+            layers (List[int]) : number of layers in each block
+            num_classes (int): number of classification classes (default 1000)
+            in_chans (int): number of input (color) channels. (default 3)
+            output_stride (int): output stride of the network, 32, 16, or 8. (default 32)
+            global_pool (str): Global pooling type. One of 'avg', 'max' (default 'avg')
+            cardinality (int): number of convolution groups for 3x3 conv in Bottleneck. (default 1)
+            base_width (int): bottleneck channels factor. `planes * base_width / 64 * cardinality` (default 64)
+            stem_width (int): number of channels in stem convolutions (default 64)
+            stem_type (str): The type of stem (default ''):
+                * '', default - a single 7x7 conv with a width of stem_width
+                * 'deep' - three 3x3 convolution layers of widths stem_width, stem_width, stem_width * 2
+                * 'deep_tiered' - three 3x3 conv layers of widths stem_width//4 * 3, stem_width, stem_width * 2
+            replace_stem_pool (bool): replace stem max-pooling layer with a 3x3 stride-2 convolution
+            block_reduce_first (int): Reduction factor for first convolution output width of residual blocks,
+                1 for all archs except senets, where 2 (default 1)
+            down_kernel_size (int): kernel size of residual block downsample path,
+                1x1 for most, 3x3 for senets (default: 1)
+            avg_down (bool): use avg pooling for projection skip connection between stages/downsample (default False)
+            act_layer (str, nn.Module): activation layer
+            norm_layer (str, nn.Module): normalization layer
+            drop_rate (float): Dropout probability before classifier, for training (default 0.)
+            zero_init_last (bool): zero-init the last weight in residual path (usually last BN affine weight)
+            block_args (dict): Extra kwargs to pass through to block module
+        """
+        super(ResNet, self).__init__()
+        block_args = block_args or dict()
+        assert output_stride in (8, 16, 32)
+        self.num_classes = num_classes
+        self.drop_rate = drop_rate
+
+        # Stem
+        deep_stem = 'deep' in stem_type
+        inplanes = stem_width * 2 if deep_stem else 64
+        if deep_stem:
+            stem_chs = (stem_width, stem_width)
+            if 'tiered' in stem_type:
+                stem_chs = (3 * (stem_width // 4), stem_width)
+            self.conv1 = nn.Sequential(*[
+                nn.Conv3d(in_chans, stem_chs[0], 3, stride=2, padding=1, bias=False),
+                norm_layer(stem_chs[0]),
+                act_layer(),
+                nn.Conv3d(stem_chs[0], stem_chs[1], 3, stride=1, padding=1, bias=False),
+                norm_layer(stem_chs[1]),
+                act_layer(),
+                nn.Conv3d(stem_chs[1], inplanes, 3, stride=1, padding=1, bias=False)])
+        else:
+            self.conv1 = nn.Conv3d(in_chans, inplanes, kernel_size=7, stride=2, padding=3, bias=False)
+        self.bn1 = norm_layer(inplanes)
+        self.act1 = act_layer()
+        self.feature_info = [dict(num_chs=inplanes, reduction=2, module='act1')]
+
+        # Stem pooling. The name 'maxpool' remains for weight compatibility.
+        if replace_stem_pool:
+            self.maxpool = nn.Sequential(*filter(None, [
+                nn.Conv3d(inplanes, inplanes, 3, stride=2, padding=1, bias=False),
+                norm_layer(inplanes),
+                act_layer(),
+            ]))
+        else:
+            self.maxpool = nn.MaxPool3d(kernel_size=3, stride=2, padding=1)
+
+        # Feature Blocks
+        block_fns = to_ntuple(len(channels))(block)
+        stage_modules, stage_feature_info = make_blocks(
+            block_fns,
+            channels,
+            layers,
+            inplanes,
+            cardinality=cardinality,
+            base_width=base_width,
+            output_stride=output_stride,
+            reduce_first=block_reduce_first,
+            avg_down=avg_down,
+            down_kernel_size=down_kernel_size,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+            **block_args,
+        )
+        for stage in stage_modules:
+            self.add_module(*stage)  # layer1, layer2, etc
+        self.feature_info.extend(stage_feature_info)
+
+        # Head (Pooling and Classifier)
+        self.num_features = self.head_hidden_size = channels[-1] * block_fns[-1].expansion
+        if global_pool == 'avg':
+            self.global_pool = nn.AdaptiveAvgPool3d((1, 1, 1))
+        elif global_pool == 'max':
+            self.global_pool = nn.AdaptiveMaxPool3d((1, 1, 1))
+        else:
+            raise NotImplementedError('Global pooling type not supported: {}'.format(global_pool))
+        
+        self.fc = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
+
+        self.init_weights(zero_init_last=zero_init_last)
+
+    @torch.jit.ignore
+    def init_weights(self, zero_init_last: bool = True):
+        for n, m in self.named_modules():
+            if isinstance(m, nn.Conv3d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+        if zero_init_last:
+            for m in self.modules():
+                if hasattr(m, 'zero_init_last'):
+                    m.zero_init_last()
+
+    @torch.jit.ignore
+    def group_matcher(self, coarse: bool = False):
+        matcher = dict(stem=r'^conv1|bn1|maxpool', blocks=r'^layer(\d+)' if coarse else r'^layer(\d+)\.(\d+)')
+        return matcher
+
+    @torch.jit.ignore
+    def get_classifier(self, name_only: bool = False):
+        return 'fc' if name_only else self.fc
+
+    def reset_classifier(self, num_classes: int, global_pool: str = 'avg'):
+        self.num_classes = num_classes
+        if global_pool == 'avg':
+            self.global_pool = nn.AdaptiveAvgPool3d((1, 1, 1))
+        elif global_pool == 'max':
+            self.global_pool = nn.AdaptiveMaxPool3d((1, 1, 1))
+        else:
+            raise NotImplementedError('Global pooling type not supported: {}'.format(global_pool))
+        
+        self.fc = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
+
+    def forward_intermediates(
+            self,
+            x: torch.Tensor,
+            indices: Optional[Union[int, List[int]]] = None,
+            norm: bool = False,
+            stop_early: bool = False,
+            output_fmt: str = 'NCHWD',
+            intermediates_only: bool = False,
+    ) -> Union[List[torch.Tensor], Tuple[torch.Tensor, List[torch.Tensor]]]:
+        """ Forward features that returns intermediates.
+
+        Args:
+            x: Input image tensor
+            indices: Take last n blocks if int, all if None, select matching indices if sequence
+            norm: Apply norm layer to compatible intermediates
+            stop_early: Stop iterating over blocks when last desired intermediate hit
+            output_fmt: Shape of intermediate feature outputs
+            intermediates_only: Only return intermediate features
+        Returns:
+
+        """
+        assert output_fmt in ('NCHWD',), 'Output shape must be NCHWD.'
+        intermediates = []
+        take_indices, max_index = feature_take_indices(5, indices)
+
+        # forward pass
+        feat_idx = 0
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.act1(x)
+        if feat_idx in take_indices:
+            intermediates.append(x)
+        x = self.maxpool(x)
+
+        layer_names = ('layer1', 'layer2', 'layer3', 'layer4')
+        if stop_early:
+            layer_names = layer_names[:max_index]
+        for n in layer_names:
+            feat_idx += 1
+            x = getattr(self, n)(x)  # won't work with torchscript, but keeps code reasonable, FML
+            if feat_idx in take_indices:
+                intermediates.append(x)
+
+        if intermediates_only:
+            return intermediates
+
+        return x, intermediates
+
+    def prune_intermediate_layers(
+            self,
+            indices: Union[int, List[int]] = 1,
+            prune_norm: bool = False,
+            prune_head: bool = True,
+    ):
+        """ Prune layers not required for specified intermediates.
+        """
+        take_indices, max_index = feature_take_indices(5, indices)
+        layer_names = ('layer1', 'layer2', 'layer3', 'layer4')
+        layer_names = layer_names[max_index:]
+        for n in layer_names:
+            setattr(self, n, nn.Identity())
+        if prune_head:
+            self.reset_classifier(0, '')
+        return take_indices
+
+    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.act1(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        return x
+
+    def forward_head(self, x: torch.Tensor, pre_logits: bool = False) -> torch.Tensor:
+        x = self.global_pool(x)
+        x = x.flatten(1)
+        if self.drop_rate:
+            x = F.dropout(x, p=float(self.drop_rate), training=self.training)
+        return x if pre_logits else self.fc(x)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.forward_features(x)
+        x = self.forward_head(x)
+        return x
+    
+    @classmethod
+    def from_pretrained(
+            cls,
+            checkpoint_path_or_url: Union[str, os.PathLike],
+            verbose: bool = True,
+            **kwargs
+    ) -> 'ResNet':
+        """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
+
+        if _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
+
+    
+
+def resnet18(
+    checkpoint_path_or_url: Optional[str] = None,
+    **kwargs
+) -> ResNet:
+    """ResNet-18 model with 3D operations.
+    """
+    kwargs = dict(
+        block=BasicBlock,
+        layers=[2, 2, 2, 2],
+        cardinality=1,
+        **kwargs,
+    )
+    if checkpoint_path_or_url:
+        return ResNet.from_pretrained(checkpoint_path_or_url, **kwargs)
+    return ResNet(**kwargs)
+
+
+def resnet34(
+    checkpoint_path_or_url: Optional[str] = None,
+    **kwargs
+) -> ResNet:
+    """ResNet-34 model with 3D operations.
+    """
+    kwargs = dict(
+        block=BasicBlock,
+        layers=[3, 4, 6, 3],
+        cardinality=1,
+        **kwargs,
+    )
+    if checkpoint_path_or_url:
+        return ResNet.from_pretrained(checkpoint_path_or_url, **kwargs)
+    return ResNet(**kwargs)
+
+
+def resnet50(
+    checkpoint_path_or_url: Optional[str] = None,
+    **kwargs
+) -> ResNet:
+    """ResNet-50 model with 3D operations.
+    """
+    kwargs = dict(
+        block=Bottleneck,
+        layers=[3, 4, 6, 3],
+        cardinality=1,
+        **kwargs,
+    )
+    if checkpoint_path_or_url:
+        return ResNet.from_pretrained(checkpoint_path_or_url, **kwargs)
+    return ResNet(**kwargs)
+
+
+def resnet101(
+    checkpoint_path_or_url: Optional[str] = None,
+    **kwargs
+) -> ResNet:
+    """ResNet-101 model with 3D operations.
+    """
+    kwargs = dict(
+        block=Bottleneck,
+        layers=[3, 4, 23, 3],
+        cardinality=1,
+        **kwargs,
+    )
+    if checkpoint_path_or_url:
+        return ResNet.from_pretrained(checkpoint_path_or_url, **kwargs)
+    return ResNet(**kwargs)
+
+
+def resnext50(
+    checkpoint_path_or_url: Optional[str] = None,
+    **kwargs
+) -> ResNet:
+    """ResNeXt-50 model with 3D operations.
+    """
+    kwargs = dict(
+        block=Bottleneck,
+        layers=[3, 4, 6, 3],
+        cardinality=32,
+        base_width=4,
+        **kwargs,
+    )
+    if checkpoint_path_or_url:
+        return ResNet.from_pretrained(checkpoint_path_or_url, **kwargs)
+    return ResNet(**kwargs)
+
+
+def resnext101(
+    checkpoint_path_or_url: Optional[str] = None,
+    **kwargs
+) -> ResNet:
+    """ResNeXt-101 model with 3D operations.
+    """
+    kwargs = dict(
+        block=Bottleneck,
+        layers=[3, 4, 23, 3],
+        cardinality=32,
+        base_width=8,
+        **kwargs,
+    )
+    if checkpoint_path_or_url:
+        return ResNet.from_pretrained(checkpoint_path_or_url, **kwargs)
+    return ResNet(**kwargs)

spectre/models/seomt.py

@@ -0,0 +1,394 @@
+# Adapted from https://github.com/tue-mps/eomt/
+
+import math
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from spectre.models import VisionTransformer
+from spectre.models.layers import LayerNorm3d
+
+
+class ScaleBlock(nn.Module):
+    def __init__(
+        self, 
+        embed_dim: int, 
+        scale_factors: Union[int, Tuple[int, int, int]] = (2, 2, 2),
+        conv1_layer: nn.Module = nn.ConvTranspose3d,
+    ):
+        super().__init__()
+
+        self.conv1 = conv1_layer(
+            embed_dim,
+            embed_dim,
+            kernel_size=scale_factors,
+            stride=scale_factors,
+        )
+        self.act = nn.GELU()
+        self.conv2 = nn.Conv3d(
+            embed_dim,
+            embed_dim,
+            kernel_size=3,
+            padding=1,
+            groups=embed_dim,
+            bias=False,
+        )
+        self.norm = LayerNorm3d(embed_dim)
+
+    def forward(self, x):
+        # print(x.shape)
+        x = self.conv1(x)
+        x = self.act(x)
+        x = self.conv2(x)
+        x = self.norm(x)
+
+        return x
+
+
+def compute_upscale_stages(patch_size, min_size=4):
+    # Compute how many times to upscale per dimension
+    num_stages = []
+    for size in patch_size:
+        stages = max(0, int(math.log2(size)) - int(math.log2(min_size)))
+        num_stages.append(stages)
+    return num_stages
+
+def voxel_shuffle_3d(x: torch.Tensor, r: int = 2) -> torch.Tensor:
+    """
+    Rearranges channels of a 5D tensor (N, C*r^3, D, H, W) to
+    (N, C, D*r, H*r, W*r).
+    """
+    n, c, d, h, w = x.size()
+    assert c % (r ** 3) == 0, f"Channels {c} not divisible by r^3={r**3}"
+    out_c = c // (r ** 3)
+    x = x.view(n, out_c, r, r, r, d, h, w)          # (N, C, r, r, r, D, H, W)
+    x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()  # (N, C, D, r, H, r, W, r)
+    x = x.view(n, out_c, d * r, h * r, w * r)       # (N, C, D*r, H*r, W*r)
+    return x
+
+
+class MLPUpBlock3D(nn.Module):
+    """
+    2x upsampling via:
+      - 1x1x1 Conv (per-voxel MLP) expanding channels by 2^3
+      - 3D voxel shuffle to double D/H/W
+      - optional norm + activation
+    """
+    def __init__(self, channels: int, norm=nn.Identity, activation=nn.GELU):
+        super().__init__()
+        self.proj = nn.Conv3d(channels, channels * 8, kernel_size=1, bias=True)
+        self.norm = norm(channels) if norm is not None else nn.Identity()
+        self.act = activation() if activation is not None else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.proj(x)           # (N, C*8, D, H, W)
+        x = voxel_shuffle_3d(x, 2) # (N, C, 2D, 2H, 2W)
+        x = self.norm(x)
+        x = self.act(x)
+        return x
+
+
+class SmolMLPDecoder3D(nn.Module):
+    """
+    Simple decoder with two MLP upsampling layers (2x each) for total 4x upsampling.
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        norm=LayerNorm3d,
+        activation=nn.GELU,
+    ):
+        super().__init__()
+        self.up1 = MLPUpBlock3D(in_channels, norm=norm, activation=activation)
+        self.up2 = MLPUpBlock3D(in_channels, norm=norm, activation=activation)
+        self.head = nn.Conv3d(in_channels, out_channels, kernel_size=1, bias=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.up1(x)   # 2x
+        x = self.up2(x)   # another 2x => 4x total
+        x = self.head(x)  # map to desired output channels
+        return x
+    
+class SpatialMLPDecoder3D(nn.Module):
+    """
+    Per-voxel MLP that expands logits from (B,Q,H,W,D) to (B,Q,H*s,W*s,D*s)
+    by predicting a learned s^3 block per voxel, then rearranging spatially.
+    """
+    def __init__(self, num_classes: int, upscale_factor: int = 4, hidden_mul: int = 4):
+        super().__init__()
+        self.num_classes = num_classes
+        self.s = upscale_factor
+        hidden = hidden_mul * num_classes
+        self.mlp = nn.Sequential(
+            nn.Linear(num_classes, hidden),
+            nn.GELU(),
+            nn.Linear(hidden, num_classes * (upscale_factor ** 3)),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: (B, Q, H, W, D)
+        B, Q, H, W, D = x.shape
+        s = self.s
+        x = x.permute(0, 2, 3, 4, 1).contiguous().view(B * H * W * D, Q)  # (BHW D, Q)
+        x = self.mlp(x)                                                   # (BHW D, Q*s^3)
+        x = x.view(B, H, W, D, Q, s, s, s)                                # (B,H,W,D,Q,s,s,s)
+        x = x.permute(0, 4, 1, 5, 2, 6, 3, 7).contiguous()                # (B,Q,H,s,W,s,D,s)
+        x = x.view(B, Q, H * s, W * s, D * s)
+        return x
+
+
+class SimpleConvDecoder3D(nn.Module):
+    """
+    Depthwise ConvTranspose3d upsampler:
+    - Assumes input & output channels == num_classes
+    - Uses groups=num_classes for class-wise deconvolution
+    """
+    def __init__(self, num_classes: int, upscale_factor: int = 4):
+        super().__init__()
+        s = upscale_factor
+        self.deconv = nn.ConvTranspose3d(
+            in_channels=num_classes,
+            out_channels=num_classes,
+            kernel_size=s,
+            stride=s,
+            padding=0,
+            output_padding=0,
+            groups=num_classes,
+            bias=True,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: (B, Q, H, W, D)
+        return self.deconv(x)
+
+
+class SEoMT(nn.Module):
+    def __init__(
+        self,
+        backbone: VisionTransformer,
+        num_classes: int,
+        # num_q: int,
+        num_blocks=4,
+        masked_attn_enabled=True,
+        return_only_final_layer=False,
+        upscale_output=True,
+        for_nnunet=False,
+        decoder=False,
+    ):
+        super().__init__()
+        self.backbone = backbone
+        self.num_q = num_classes
+        self.num_blocks = num_blocks
+        self.masked_attn_enabled = masked_attn_enabled
+        self.return_only_final_layer = return_only_final_layer
+        self.upscale_output = upscale_output
+        self.register_buffer("attn_mask_probs", torch.ones(num_blocks))
+        self.for_nnunet = for_nnunet
+        self.q = nn.Embedding(num_classes, self.backbone.embed_dim)
+
+
+        self.mask_head = nn.Sequential(
+            nn.Linear(self.backbone.embed_dim, self.backbone.embed_dim),
+            nn.GELU(),
+            nn.Linear(self.backbone.embed_dim, self.backbone.embed_dim),
+            nn.GELU(),
+            nn.Linear(self.backbone.embed_dim, self.backbone.embed_dim),
+        )
+
+        patch_size = self.backbone.patch_embed.patch_size
+        num_upscale_stages = compute_upscale_stages(patch_size, min_size=4)
+
+        # Build per-stage scale factors list
+        max_stages = max(num_upscale_stages)
+        upscale_blocks = []
+        for stage_idx in range(max_stages):
+            # for each dimension, upscale by 2 only if this dimension still has
+            # remaining upscales at this stage
+            scale_factors = tuple(
+                2 if stage_idx < num_upscale_stages[dim] else 1
+                for dim in range(len(patch_size))
+            )
+            upscale_blocks.append(ScaleBlock(self.backbone.embed_dim,
+                                            scale_factors=scale_factors))
+
+        self.upscale = nn.Sequential(*upscale_blocks)
+
+    def _predict(self, x: torch.Tensor, stage: int = None):
+        q = x[:, : self.num_q, :]
+        # print(stage)
+        # class_logits = self.class_head(q)
+        x = x[:, self.num_q + self.backbone.num_prefix_tokens :, :]
+        x = x.transpose(1, 2).reshape(
+            x.shape[0], -1, *self.backbone.patch_embed.grid_size
+        )
+        mask_logits = torch.einsum(
+            "bqc, bchwd -> bqhwd", self.mask_head(q), self.upscale(x)
+        )
+
+
+        return mask_logits
+
+    @torch.compiler.disable
+    def _disable_attn_mask(self, attn_mask, prob):
+        if prob < 1:
+            random_queries = (
+                torch.rand(attn_mask.shape[0], self.num_q, device=attn_mask.device)
+                > prob
+            )
+            attn_mask[
+                :, : self.num_q, self.num_q + self.backbone.num_prefix_tokens :
+            ][random_queries] = True
+
+        return attn_mask
+
+    def _attn(self, module: 'Attention', x: torch.Tensor, mask: Optional[torch.Tensor], rope = None):
+        B, N, C = x.shape
+
+        q = module.q(x).reshape(B, N, module.num_heads, module.head_dim).permute(0, 2, 1, 3)
+        kv = module.kv(x).reshape(B, N, 2, module.num_heads, module.head_dim)
+        k, v = kv.permute(2, 0, 3, 1, 4).unbind(0)
+        q, k = module.q_norm(q), module.k_norm(k)
+
+        if mask is not None:
+            mask = mask[:, None, ...].expand(-1, module.num_heads, -1, -1)
+
+        dropout_p = module.attn_drop.p if self.training else 0.0
+
+        if rope is not None:
+            if isinstance(rope, list):
+                rope = tuple(torch.stack([r[i] for r in rope], dim=0) for i in range(2))
+            q, k = module.apply_rotary_pos_emb(q, k, rope)
+
+        if module.fused_attn:
+            x = F.scaled_dot_product_attention(q, k, v, mask, dropout_p)
+        else:
+            attn = (q @ k.transpose(-2, -1)) * module.scale
+            if mask is not None:
+                attn = attn.masked_fill(~mask, float("-inf"))
+            attn = F.softmax(attn, dim=-1)
+            attn = module.attn_drop(attn)
+            x = attn @ v
+
+        x = module.proj_drop(module.proj(x.transpose(1, 2).reshape(B, N, C)))
+
+        return x
+
+    def forward(self, x: torch.Tensor):
+
+        if self.for_nnunet: # swap data order, will be incoming at czyx - cxyz
+            x = x.permute(0, 1, 3, 4, 2).contiguous()
+
+        self.backbone.patch_embed.set_input_size(x.shape[2:])
+        x = self.backbone.patch_embed(x)
+        x, rope = self.backbone._pos_embed(x)
+        x = self.backbone.patch_drop(x)
+        x = self.backbone.norm_pre(x)
+        attn_mask = None
+        mask_logits_per_layer = []
+
+        for i, block in enumerate(self.backbone.blocks):
+            if i == len(self.backbone.blocks) - self.num_blocks:
+                x = torch.cat(
+                    (self.q.weight[None, :, :].expand(x.shape[0], -1, -1), x), dim=1
+                )
+
+            if (
+                self.masked_attn_enabled
+                and i >= len(self.backbone.blocks) - self.num_blocks
+            ):
+                mask_logits = self._predict(self.backbone.norm(x))
+
+                if self.for_nnunet:
+                     # swap back to czyx
+                    
+                    if self.upscale_output:
+                        # Upscale to original input size / stage
+
+                        stage = len(self.backbone.blocks) -i 
+                        if stage is not None:
+                            input_size = tuple(
+                                int(self.backbone.patch_embed.patch_size[dim] * self.backbone.patch_embed.grid_size[dim] / 2**(stage))
+                                for dim in range(len(self.backbone.patch_embed.patch_size))
+                            )
+
+                        mask_logits_per_layer.append(F.interpolate(mask_logits, input_size, mode="trilinear").permute(0, 1, 4, 2, 3).contiguous())
+                    else:
+                        mask_logits_per_layer.append(mask_logits.permute(0, 1, 4, 2, 3).contiguous())
+                else:
+                    mask_logits_per_layer.append(mask_logits)
+
+                attn_mask = torch.ones(
+                    x.shape[0],
+                    x.shape[1],
+                    x.shape[1],
+                    dtype=torch.bool,
+                    device=x.device,
+                )
+                interpolated = F.interpolate(
+                    mask_logits,
+                    self.backbone.patch_embed.grid_size,
+                    mode="trilinear",
+                )
+                interpolated = interpolated.view(
+                    interpolated.size(0), interpolated.size(1), -1
+                )
+                attn_mask[
+                    :,
+                    : self.num_q,
+                    self.num_q + self.backbone.num_prefix_tokens :,
+                ] = (
+                    interpolated > 0
+                )
+                attn_mask = self._disable_attn_mask(
+                    attn_mask,
+                    self.attn_mask_probs[
+                        i - len(self.backbone.blocks) + self.num_blocks
+                    ],
+                )
+            x = x + block.drop_path1(
+                block.ls1(self._attn(block.attn, block.norm1(x), attn_mask, rope=rope))
+            )
+            x = x + block.drop_path2(block.ls2(block.mlp(block.norm2(x))))
+
+        mask_logits = self._predict(self.backbone.norm(x))
+        if self.for_nnunet:
+            input_size = tuple(
+                                int(self.backbone.patch_embed.patch_size[dim] * self.backbone.patch_embed.grid_size[dim] / 2**0)
+                                for dim in range(len(self.backbone.patch_embed.patch_size))
+                            )
+            mask_logits_per_layer.append(F.interpolate(mask_logits, input_size, mode="trilinear").permute(0, 1, 4, 2, 3).contiguous())
+        else:
+            mask_logits_per_layer.append(mask_logits)
+
+        if self.for_nnunet:
+            # return in reversed order for deep supervision
+            mask_logits_per_layer = mask_logits_per_layer[::-1]
+        return mask_logits_per_layer if not self.return_only_final_layer else mask_logits_per_layer[0]
+    
+
+if __name__ == "__main__":
+    from spectre.models import vit_large_patch16_128
+
+    model = SEoMT(
+        backbone=vit_large_patch16_128(pos_embed='rope',
+            rope_kwargs={
+                "base": 1000.0,  # works for most 3D models
+            },),
+        num_classes=4,
+        num_blocks=4,
+        masked_attn_enabled=True,
+        return_only_final_layer=True,
+        for_nnunet=True,
+        upscale_output=True,
+        decoder=False,
+    )
+    # print number of parameters
+    print(f"Number of parameters: {sum(p.numel() for p in model.parameters() if p.requires_grad)}")
+
+    x = torch.randn(2, 1, 64, 128, 128)
+    out = model(x)
+    for o in out:
+        print(o.shape)

spectre/models/upsample_anything.py

@@ -0,0 +1,319 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.optim.lr_scheduler import LambdaLR
+
+
+def UPA(hr_image, lr_volume, device="cuda", use_amp=True):
+    """
+    hr_image: numpy or torch [C,Hh,Wh,Dh]
+    lr_volume: torch [1,C,Hl,Wl,Dl]
+    """
+
+    hr = torch.as_tensor(hr_image).unsqueeze(0).float().to(device)
+
+    _, _, Hh, Wh, Dh = hr.shape
+    _, _, Hl, Wl, Dl = lr_volume.shape
+    scale = Hh // Hl
+    assert Wh // Wl == scale and Dh // Dl == scale, "Inconsistent scale factors"
+
+    lr_volume = lr_volume.to(device).float()
+    lr = F.interpolate(hr, scale_factor=1/scale, mode="trilinear", align_corners=False)
+
+    model = LearnablePixelwiseAnisoJBU3D(
+        Hl, Wl, Dl, scale=scale
+    ).to(device)
+
+    model.train()
+    opt = torch.optim.Adam(model.parameters(), lr=1e-1)
+    max_steps = 350
+    gamma = (1e-9 / 1e-1) ** (1.0 / max_steps)
+    scheduler = LambdaLR(opt, lr_lambda=lambda step: gamma ** step)
+    scaler = torch.amp.GradScaler(device=device, enabled=use_amp)
+
+    for step in range(max_steps):
+        opt.zero_grad(set_to_none=True)
+        with torch.amp.autocast(device_type=device, enabled=use_amp):
+            pred = model(lr, hr)
+            loss = F.l1_loss(pred, hr)
+
+        scaler.scale(loss).backward()
+        scaler.step(opt)
+        scaler.update()
+        scheduler.step()
+
+        if step % 50 == 0:
+            print(f"step {step}: loss={loss.item():.5f}")
+
+    model.eval()
+    with torch.inference_mode(), \
+        torch.amp.autocast(device_type=device, enabled=use_amp, dtype=torch.float16):
+        out = model(lr_volume, hr)
+ 
+    return out
+
+
+@torch.no_grad()
+def _build_offsets_3d(R_max: int, device):
+    offs = torch.arange(-R_max, R_max + 1, device=device)
+    dX, dY, dZ = torch.meshgrid(offs, offs, offs, indexing="ij")
+    return (
+        dX.reshape(-1),
+        dY.reshape(-1),
+        dZ.reshape(-1),
+    )  # [K]
+
+
+def gather_lr_scalar_3d(map_lr, Ui, Vi, Wi):
+    """
+    map_lr: [1,1,Hl,Wl,Dl] or [Hl,Wl,Dl]
+    Ui,Vi,Wi: [Bn,Hh,Wh,Dh]
+    """
+    Hl, Wl, Dl = map_lr.shape[-3:]
+    flat = Hl * Wl * Dl
+    idx = (Ui * Wl * Dl + Vi * Dl + Wi).reshape(-1)
+    t = map_lr.view(flat)
+    vals = t.index_select(0, idx)
+    return vals.view(Ui.shape)
+
+
+def gs_jbu_aniso_noparent_3d(
+    feat_lr,          # [1,C,Hl,Wl,Dl]
+    guide_hr,         # [1,G,Hh,Wh,Dh]
+    scale,
+    sigma_x_map,
+    sigma_y_map,
+    sigma_z_map,
+    sigma_r_map,
+    R_max=3,
+    alpha_dyn=2.0,
+    C_chunk=64,
+    Nn_chunk=125,
+):
+    _, C, Hl, Wl, Dl = feat_lr.shape
+    _, _, Hh, Wh, Dh = guide_hr.shape
+    device = feat_lr.device
+    dtype_feat = feat_lr.dtype
+
+    # HR grid
+    x = torch.arange(Hh, device=device, dtype=torch.float32)
+    y = torch.arange(Wh, device=device, dtype=torch.float32)
+    z = torch.arange(Dh, device=device, dtype=torch.float32)
+    X, Y, Z = torch.meshgrid(x, y, z, indexing="ij")
+
+    u = (X + 0.5) / scale - 0.5
+    v = (Y + 0.5) / scale - 0.5
+    w = (Z + 0.5) / scale - 0.5
+
+    uc = torch.round(u).clamp(0, Hl - 1).long()
+    vc = torch.round(v).clamp(0, Wl - 1).long()
+    wc = torch.round(w).clamp(0, Dl - 1).long()
+
+    # Dynamic radius
+    sigma_eff = torch.maximum(
+        sigma_x_map,
+        torch.maximum(sigma_y_map, sigma_z_map),
+    )
+    sigma_eff_hr = F.interpolate(
+        sigma_eff, (Hh, Wh, Dh), mode="trilinear", align_corners=False
+    )
+    # sigma_eff_hr = sigma_eff_hr.squeeze(0).squeeze(0)
+    R_map = torch.ceil(alpha_dyn * sigma_eff_hr).clamp(1, R_max).long()
+
+    dX_all, dY_all, dZ_all = _build_offsets_3d(R_max, device)
+
+    num = torch.zeros(C, Hh, Wh, Dh, device=device, dtype=torch.float32)
+    den = torch.zeros(Hh, Wh, Dh, device=device, dtype=torch.float32)
+    m = torch.full((Hh, Wh, Dh), -1e9, device=device, dtype=torch.float32)
+
+    feat_flat = feat_lr[0].permute(1, 2, 3, 0).reshape(-1, C)
+    guide_lr = F.interpolate(
+        guide_hr, (Hl, Wl, Dl), mode="trilinear", align_corners=False
+    )
+
+    for n0 in range(0, len(dX_all), Nn_chunk):
+        dX = dX_all[n0:n0+Nn_chunk][:, None, None, None]
+        dY = dY_all[n0:n0+Nn_chunk][:, None, None, None]
+        dZ = dZ_all[n0:n0+Nn_chunk][:, None, None, None]
+
+        Ui = torch.clamp(uc.unsqueeze(0) + dX, 0, Hl - 1)
+        Vi = torch.clamp(vc.unsqueeze(0) + dY, 0, Wl - 1)
+        Wi = torch.clamp(wc.unsqueeze(0) + dZ, 0, Dl - 1)
+
+        # mask = (dX**2 + dY**2 + dZ**2 <= R_map[None, ...] ** 2)
+        mask = (dX**2 + dY**2 + dZ**2 <= R_map**2).squeeze(0).squeeze(0)
+
+        cx = (Ui.float() + 0.5) * scale - 0.5
+        cy = (Vi.float() + 0.5) * scale - 0.5
+        cz = (Wi.float() + 0.5) * scale - 0.5
+
+        dx = X.unsqueeze(0) - cx
+        dy = Y.unsqueeze(0) - cy
+        dz = Z.unsqueeze(0) - cz
+
+        sx = gather_lr_scalar_3d(sigma_x_map, Ui, Vi, Wi).clamp_min(1e-6)
+        sy = gather_lr_scalar_3d(sigma_y_map, Ui, Vi, Wi).clamp_min(1e-6)
+        sz = gather_lr_scalar_3d(sigma_z_map, Ui, Vi, Wi).clamp_min(1e-6)
+        sr = gather_lr_scalar_3d(sigma_r_map, Ui, Vi, Wi).clamp_min(1e-6)
+
+        log_ws = (
+            -(dx**2)/(2*sx**2)
+            -(dy**2)/(2*sy**2)
+            -(dz**2)/(2*sz**2)
+        )
+
+        diff2 = 0.0
+        for g in range(guide_hr.shape[1]):
+            g0 = gather_lr_scalar_3d(guide_lr[0, g], Ui, Vi, Wi)
+            diff2 += (guide_hr[0, g] - g0) ** 2
+
+        log_wr = -diff2 / (2 * sr**2 + 1e-8)
+        log_w = torch.where(mask, log_ws + log_wr, -1e9)
+
+        m_chunk = log_w.max(dim=0).values
+        m_new = torch.maximum(m, m_chunk)
+
+        scale_old = torch.exp(m - m_new)
+        num *= scale_old
+        den *= scale_old
+
+        w = torch.exp(log_w - m_new)
+        den += w.sum(0)
+
+        idx_flat = (Ui * Wl * Dl + Vi * Dl + Wi).reshape(-1)
+
+        for c0 in range(0, C, C_chunk):
+            c1 = min(c0 + C_chunk, C)
+            f = feat_flat.index_select(0, idx_flat)[:, c0:c1]
+            f = f.view(w.shape + (c1 - c0,))
+            num[c0:c1] += (f * w[..., None]).sum(0).permute(3, 0, 1, 2)
+
+        m = m_new
+
+    out = (num / den.clamp_min(1e-8)).unsqueeze(0)
+    return out.to(dtype_feat)
+
+
+class LearnablePixelwiseAnisoJBU3D(nn.Module):
+    def __init__(
+        self,
+        Hl,
+        Wl,
+        Dl,
+        scale,
+        init_sigma=1.5,
+        init_sigma_r=0.1,
+        R_max=3,
+        alpha_dyn=2.0,
+    ):
+        super().__init__()
+        self.scale = scale
+        self.R_max = R_max
+        self.alpha_dyn = alpha_dyn
+
+        self.sx_raw = nn.Parameter(torch.full((1,1,Hl,Wl,Dl), math.log(init_sigma)))
+        self.sy_raw = nn.Parameter(torch.full((1,1,Hl,Wl,Dl), math.log(init_sigma)))
+        self.sz_raw = nn.Parameter(torch.full((1,1,Hl,Wl,Dl), math.log(init_sigma)))
+        self.sr_raw = nn.Parameter(torch.full((1,1,Hl,Wl,Dl), math.log(init_sigma_r)))
+
+    def forward(self, feat_lr, guide_hr):
+        return gs_jbu_aniso_noparent_3d(
+            feat_lr,
+            guide_hr,
+            self.scale,
+            torch.exp(self.sx_raw),
+            torch.exp(self.sy_raw),
+            torch.exp(self.sz_raw),
+            torch.exp(self.sr_raw),
+            R_max=self.R_max,
+            alpha_dyn=self.alpha_dyn,
+        )
+
+
+if __name__ == "__main__":
+    import argparse
+
+    import numpy as np
+    import nibabel as nib
+    import monai.transforms as transforms
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--image_path", type=str, required=True)
+    parser.add_argument("--mask_path", type=str, required=True)
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--use_amp", action="store_true")
+    args = parser.parse_args()
+
+    transform = transforms.Compose([
+        transforms.LoadImaged(keys=("image", "mask")),
+        transforms.EnsureChannelFirstd(keys=("image", "mask"), channel_dim="no_channel"),
+        transforms.ScaleIntensityRanged(
+            keys=("image",),
+            a_min=-150,
+            a_max=250,
+            b_min=0.0,
+            b_max=1.0,
+            clip=True,
+        ),
+        transforms.Orientationd(keys=("image", "mask"), axcodes="RAS"),
+        transforms.RandWeightedCropd(
+            keys=("image", "mask"), 
+            w_key="mask", 
+            spatial_size=(128, 128, 64), 
+            num_samples=1,
+        ),
+        transforms.CopyItemsd(keys=("mask"), times=1, names=("mask_low_res")),
+        transforms.Resized(keys=("mask_low_res"), spatial_size=(16, 16, 8), mode="nearest", align_corners=False)
+    ])
+    sample = transform({
+        "image": args.image_path,
+        "mask": args.mask_path,
+    })[0]
+
+    nib.save(
+        nib.Nifti1Image(
+            (F.interpolate(sample["mask_low_res"].unsqueeze(0), size=(128, 128, 64), mode="nearest").squeeze(0).squeeze(0).cpu().numpy().astype(np.uint8)),
+            affine=np.eye(4),
+        ),
+        "mask_low_res_upscaled.nii.gz",
+    )
+
+    sample["mask_low_res"] = F.one_hot(
+        sample["mask_low_res"].long().squeeze(0), num_classes=4,
+    ).permute(3, 0, 1, 2).unsqueeze(0).float()
+
+    print(sample["mask_low_res"].shape)
+
+    mask_out = UPA(
+        sample["image"],
+        sample["mask_low_res"],
+        device=args.device,
+        use_amp=args.use_amp,
+    )
+
+    mask_out = mask_out.argmax(dim=1, keepdim=True)
+
+    nib.save(
+        nib.Nifti1Image(
+            (sample["image"] * 255).squeeze(0).cpu().numpy().astype(np.uint8),
+            affine=np.eye(4),
+        ),
+        "image.nii.gz",
+    )
+    nib.save(
+        nib.Nifti1Image(
+            sample["mask"].squeeze(0).cpu().numpy().astype(np.uint8),
+            affine=np.eye(4),
+        ),
+        "mask.nii.gz",
+    )
+    torch.save(mask_out.squeeze(0).squeeze(0).cpu(), "upsampled_mask.pt")
+    nib.save(
+        nib.Nifti1Image(
+            mask_out.squeeze(0).squeeze(0).cpu().numpy().astype(np.uint8),
+            affine=np.eye(4),
+        ),
+        "upsampled_mask.nii.gz",
+    )
+    

spectre/models/vision_transformer.py

@@ -0,0 +1,835 @@
+import os
+from functools import partial
+from urllib.parse import urlparse
+from typing import (
+    Tuple, Union, Callable, Literal, 
+    Optional, Type, Set, List, Dict, Any,
+)
+
+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, 
+    RotaryPositionEmbedding,
+)
+from spectre.utils import (
+    resample_abs_pos_embed, 
+    feature_take_indices, 
+    global_pool_nlc,
+)
+
+
+class Block(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        attn_mode: str = 'mha',
+        q_proj_dim: Optional[int] = None,
+        kv_proj_dim: Optional[int] = None,
+        mlp_ratio: float = 4.,
+        qkv_bias: bool = False,
+        qk_norm: bool = False,
+        proj_bias: bool = True,
+        proj_drop: float = 0.,
+        attn_drop: float = 0.,
+        init_values: Optional[float] = None,
+        drop_path: float = 0.,
+        act_layer: Type[nn.Module] = nn.GELU,
+        norm_layer: Type[nn.Module] = nn.LayerNorm,
+        mlp_layer: Type[nn.Module] = Mlp,
+    ) -> None:
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim,
+            num_heads=num_heads,
+            mode=attn_mode,
+            q_proj_dim=q_proj_dim,
+            kv_proj_dim=kv_proj_dim,
+            qkv_bias=qkv_bias,
+            qk_norm=qk_norm,
+            proj_bias=proj_bias,
+            attn_drop=attn_drop,
+            proj_drop=proj_drop,
+            norm_layer=norm_layer,
+        )
+        self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+        self.norm2 = norm_layer(dim)
+        self.mlp = mlp_layer(
+            in_features=dim,
+            hidden_features=int(dim * mlp_ratio),
+            act_layer=act_layer,
+            bias=proj_bias,
+            drop=proj_drop,
+        )
+        self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path2 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+    def forward(
+        self, 
+        x: torch.Tensor, 
+        rope = None
+    ) -> torch.Tensor:
+        x = x + self.drop_path1(self.ls1(self.attn(self.norm1(x), rope=rope)))
+        x = x + self.drop_path2(self.ls2(self.mlp(self.norm2(x))))
+        return x
+
+
+class VisionTransformer(nn.Module):
+    """ Vision Transformer with 3D Patch Embedding
+    """
+    def __init__(
+        self, 
+        img_size: Union[int, Tuple[int, int, int]] = (128, 128, 64),
+        patch_size: Union[int, Tuple[int, int, int]] = (16, 16, 8),
+        in_chans: int = 1,
+        num_classes: int = 1000,
+        global_pool: Literal['', 'avg', 'avgmax', 'max', 'token', 'map'] = 'token',
+        embed_dim: int = 768,
+        depth: int = 12,
+        num_heads: int = 12,
+        attn_mode: str = 'mha',
+        q_proj_dim: Optional[int] = None,
+        kv_proj_dim: Optional[int] = None,
+        mlp_ratio: float = 4.,
+        qkv_bias: bool = True,
+        qk_norm: bool = False,
+        proj_bias: bool = True,
+        init_values: Optional[float] = None,
+        class_token: bool = True,
+        pos_embed: str = 'learn',
+        no_embed_class: bool = False,
+        rope_kwargs: Optional[dict] = None,
+        reg_tokens: int = 0,
+        pre_norm: bool = False,
+        final_norm: bool = True,
+        fc_norm: Optional[bool] = None,
+        dynamic_img_size: bool = False,
+        dynamic_img_pad: bool = False,
+        drop_rate: float = 0.,
+        pos_drop_rate: float = 0.,
+        patch_drop_rate: float = 0.,
+        proj_drop_rate: float = 0.,
+        attn_drop_rate: float = 0.,
+        drop_path_rate: float = 0.,
+        embed_layer: Callable = PatchEmbed,
+        embed_norm_layer: Optional[Union[Callable, Type[torch.nn.Module]]] = None,
+        norm_layer: Optional[Union[Callable, Type[torch.nn.Module]]] = None,
+        act_layer: Optional[Union[Callable, Type[torch.nn.Module]]] = None,
+        block_fn: Type[nn.Module] = Block,
+        mlp_layer: Type[nn.Module] = Mlp,
+    ) -> None:
+        """
+        Args:
+            img_size: Input image size.
+            patch_size: Patch size.
+            in_chans: Number of image input channels.
+            num_classes: Number of classes for classification head.
+            global_pool: Type of global pooling for final sequence (default: 'token').
+            embed_dim: Transformer embedding dimension.
+            depth: Depth of transformer.
+            num_heads: Number of attention heads.
+            attn_mode: Attention mode ('mha', 'mqa', 'mla').
+            q_proj_dim: Query projection dimension for 'mla' mode.
+            kv_proj_dim: Key, value projection dimension for 'mla' mode.
+            mlp_ratio: Ratio of mlp hidden dim to embedding dim.
+            qkv_bias: Enable bias for qkv projections if True.
+            init_values: Layer-scale init values (layer-scale enabled if not None).
+            class_token: Use class token.
+            pos_embed: Type of position embedding to use (default: 'learn').
+            no_embed_class: Don't include position embeddings for class (or reg) tokens for learnable pos_embed.
+            rope_kwargs: Additional arguments for rotary position embedding.
+            reg_tokens: Number of register tokens.
+            pre_norm: Enable norm after embeddings, before transformer blocks (standard in CLIP ViT).
+            final_norm: Enable norm after transformer blocks, before head (standard in most ViT).
+            fc_norm: Move final norm after pool (instead of before), if None, enabled when global_pool == 'avg'.
+            drop_rate: Head dropout rate.
+            pos_drop_rate: Position embedding dropout rate.
+            attn_drop_rate: Attention dropout rate.
+            drop_path_rate: Stochastic depth rate.
+            weight_init: Weight initialization scheme.
+            fix_init: Apply weight initialization fix (scaling w/ layer index).
+            embed_layer: Patch embedding layer.
+            embed_norm_layer: Normalization layer to use / override in patch embed module.
+            norm_layer: Normalization layer.
+            act_layer: MLP activation layer.
+            block_fn: Transformer block layer.
+        """
+        super().__init__()
+        assert global_pool in ('', 'avg', 'avgmax', 'max', 'token', 'map')
+        assert class_token or global_pool != 'token'
+        assert pos_embed in ('', 'none', 'learn', 'rope')
+        assert attn_mode in ('mha', 'mqa', 'mla')
+        rope_kwargs = {} if rope_kwargs is None else dict(rope_kwargs)
+        rope_kwargs.setdefault("dtype", torch.float32)  # robust with mixed-precision
+        use_fc_norm = global_pool in ('avg', 'avgmax', 'max') if fc_norm is None else fc_norm
+        norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
+        embed_norm_layer = embed_norm_layer
+        act_layer = act_layer or nn.GELU
+
+        self.num_classes = num_classes
+        self.global_pool = global_pool
+        self.num_features = self.head_hidden_size = self.embed_dim = embed_dim  # for consistency with other models
+        self.num_prefix_tokens = 1 if class_token else 0
+        self.num_prefix_tokens += reg_tokens
+        self.num_reg_tokens = reg_tokens
+        self.has_class_token = class_token
+        self.no_embed_class = no_embed_class  # don't embed prefix positions (includes reg)
+        self.dynamic_img_size = dynamic_img_size
+
+        embed_args = {}
+        if self.dynamic_img_size:
+            # flatten deferred until after pos embed
+            embed_args.update(dict(strict_img_size=False, output_fmt="NHWDC"))
+        elif pos_embed == 'rope':
+            embed_args['output_fmt'] = "NHWDC"
+        if embed_norm_layer is not None:
+            embed_args['norm_layer'] = embed_norm_layer
+        self.patch_embed = embed_layer(
+            img_size=img_size,
+            patch_size=patch_size,
+            in_chans=in_chans,
+            embed_dim=embed_dim,
+            bias=not pre_norm,  # disable bias if pre-norm is used (e.g. CLIP)
+            dynamic_img_pad=dynamic_img_pad,
+            **embed_args,
+        )
+        num_patches = self.patch_embed.num_patches
+        reduction = self.patch_embed.feat_ratio() if hasattr(self.patch_embed, 'feat_ratio') else patch_size
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) if class_token else None
+        self.reg_token = nn.Parameter(torch.zeros(1, reg_tokens, embed_dim)) if reg_tokens else None
+        embed_len = num_patches if no_embed_class else num_patches + self.num_prefix_tokens
+        self.pos_embed, self.rope, self.requires_per_sample_rope = None, None, False
+        if pos_embed == 'learn':
+            self.pos_embed = nn.Parameter(torch.randn(1, embed_len, embed_dim) * .02)
+        if pos_embed == 'rope':
+            self.rope = RotaryPositionEmbedding(
+                embed_dim=embed_dim,
+                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(
+                patch_drop_rate,
+                num_prefix_tokens=self.num_prefix_tokens,
+            )
+        else:
+            self.patch_drop = nn.Identity()
+        self.norm_pre = norm_layer(embed_dim) if pre_norm else nn.Identity()
+
+        dpr = [drop_path_rate * i / (depth - 1) if depth > 1 else 0.0 for i in range(depth)]  # stochastic depth decay rule
+        self.blocks = nn.Sequential(*[
+            block_fn(
+                dim=embed_dim,
+                num_heads=num_heads,
+                attn_mode=attn_mode,
+                q_proj_dim=q_proj_dim,
+                kv_proj_dim=kv_proj_dim,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                qk_norm=qk_norm,
+                proj_bias=proj_bias,
+                init_values=init_values,
+                proj_drop=proj_drop_rate,
+                attn_drop=attn_drop_rate,
+                drop_path=dpr[i],
+                norm_layer=norm_layer,
+                act_layer=act_layer,
+                mlp_layer=mlp_layer,
+            )
+            for i in range(depth)])
+        self.feature_info = [
+            dict(module=f'blocks.{i}', num_chs=embed_dim, reduction=reduction) for i in range(depth)]
+        self.norm = norm_layer(embed_dim) if final_norm and not use_fc_norm else nn.Identity()
+
+        # Classifier Head
+        if global_pool == 'map':
+            self.attn_pool = AttentionPoolLatent(
+                self.embed_dim,
+                num_heads=num_heads,
+                mlp_ratio=mlp_ratio,
+                norm_layer=norm_layer,
+                act_layer=act_layer,
+            )
+        else:
+            self.attn_pool = None
+        self.fc_norm = norm_layer(embed_dim) if final_norm and use_fc_norm else nn.Identity()
+        self.head_drop = nn.Dropout(drop_rate)
+        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+        self.init_weights()
+
+    def init_weights(self) -> None:
+        if self.pos_embed is not None and not self.pos_embed.is_meta:
+            nn.init.trunc_normal_(self.pos_embed, std=.02)
+        if self.cls_token is not None and not self.cls_token.is_meta:
+            nn.init.normal_(self.cls_token, std=1e-6)
+        if self.reg_token is not None and not self.reg_token.is_meta:
+            nn.init.normal_(self.reg_token, std=1e-6)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m: nn.Module) -> None:
+        # this fn left here for compat with downstream users
+        if isinstance(m, nn.Linear):
+            if not m.weight.is_meta:
+                nn.init.trunc_normal_(m.weight, std=.02)
+            if m.bias is not None and not m.bias.is_meta:
+                nn.init.zeros_(m.bias)
+
+    @torch.jit.ignore
+    def no_weight_decay(self) -> Set:
+        return {'pos_embed', 'cls_token', 'dist_token'}
+
+    @torch.jit.ignore
+    def get_classifier(self) -> nn.Module:
+        return self.head
+
+    def reset_classifier(self, num_classes: int, global_pool: Optional[str] = None):
+        self.num_classes = num_classes
+        if global_pool is not None:
+            assert global_pool in ('', 'avg', 'avgmax', 'max', 'token', 'map')
+            if global_pool == 'map' and self.attn_pool is None:
+                assert False, "Cannot currently add attention pooling in reset_classifier()."
+            elif global_pool != 'map' and self.attn_pool is not None:
+                self.attn_pool = None  # remove attention pooling
+            self.global_pool = global_pool
+        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+    def set_input_size(
+        self,
+        img_size: Optional[Tuple[int, int, int]] = None,
+        patch_size: Optional[Tuple[int, int, int]] = None,
+    ):
+        """Method updates the input image resolution, patch size
+
+        Args:
+            img_size: New input resolution, if None current resolution is used
+            patch_size: New patch size, if None existing patch size is used
+        """
+        prev_grid_size = self.patch_embed.grid_size
+        self.patch_embed.set_input_size(img_size=img_size, patch_size=patch_size)
+        if self.pos_embed is not None:
+            num_prefix_tokens = 0 if self.no_embed_class else self.num_prefix_tokens
+            num_new_tokens = self.patch_embed.num_patches + num_prefix_tokens
+            if num_new_tokens != self.pos_embed.shape[1]:
+                self.pos_embed = nn.Parameter(resample_abs_pos_embed(
+                    self.pos_embed,
+                    new_size=self.patch_embed.grid_size,
+                    old_size=prev_grid_size,
+                    num_prefix_tokens=num_prefix_tokens,
+                    verbose=True,
+                ))
+
+    def _pos_embed(self, x: torch.Tensor):
+        if self.pos_embed is None and self.rope is None:
+            x = x.view(x.shape[0], -1, x.shape[-1])
+            if self.reg_token is not None:
+                x = torch.cat([self.reg_token.expand(x.shape[0], -1, -1), x], dim=1)
+            if self.cls_token is not None:
+                x = torch.cat([self.cls_token.expand(x.shape[0], -1, -1), x], dim=1)
+            return x, None
+        
+        if self.dynamic_img_size or self.rope is not None:
+            B, H, W, D, C = x.shape
+            x = x.view(B, -1, C)
+
+        pos_embed, rope = None, None
+        if self.pos_embed is not None:
+            if self.dynamic_img_size:
+                prev_grid_size = self.patch_embed.grid_size
+                pos_embed = resample_abs_pos_embed(
+                    self.pos_embed,
+                    new_size=(H, W, D),
+                    old_size=prev_grid_size,
+                    num_prefix_tokens=0 if self.no_embed_class else self.num_prefix_tokens,
+                )
+            else:
+                pos_embed = self.pos_embed
+        
+        if self.rope is not None:
+            if self.requires_per_sample_rope:
+                rope = [self.rope(H=H, W=W, D=D) for _ in range(B)]
+            else:
+                rope = self.rope(H=H, W=W, D=D)
+
+        to_cat = []
+        if self.cls_token is not None:
+            to_cat.append(self.cls_token.expand(x.shape[0], -1, -1))
+        if self.reg_token is not None:
+            to_cat.append(self.reg_token.expand(x.shape[0], -1, -1))
+
+        if self.no_embed_class:
+            # deit-3, updated JAX (big vision)
+            # position embedding does not overlap with class token, add then concat
+            if pos_embed is not None:
+                x = x + pos_embed
+            if to_cat:
+                x = torch.cat(to_cat + [x], dim=1)
+        else:
+            # original timm, JAX, and deit vit impl
+            # pos_embed has entry for class token, concat then add
+            if to_cat:
+                x = torch.cat(to_cat + [x], dim=1)
+            if pos_embed is not None:
+                x = x + pos_embed
+
+        return self.pos_drop(x), rope
+
+    def forward_intermediates(
+        self,
+        x: torch.Tensor,
+        indices: Optional[Union[int, List[int]]] = None,
+        return_prefix_tokens: bool = False,
+        norm: bool = False,
+        stop_early: bool = False,
+        output_fmt: str = 'NCHWD',
+        intermediates_only: bool = False,
+        output_dict: bool = False,
+    ) -> Union[List[torch.Tensor], Tuple[torch.Tensor, List[torch.Tensor]], Dict[str, Any]]:
+        """ Forward features that returns intermediates.
+
+        Args:
+            x: Input image tensor
+            indices: Take last n blocks if int, all if None, select matching indices if sequence
+            return_prefix_tokens: Return both prefix and spatial intermediate tokens
+            norm: Apply norm layer to all intermediates
+            stop_early: Stop iterating over blocks when last desired intermediate hit
+            output_fmt: Shape of intermediate feature outputs
+            intermediates_only: Only return intermediate features
+            output_dict: Return outputs as a dictionary with 'image_features' and 'image_intermediates' keys
+        Returns:
+            A tuple with (final_features, intermediates), a list of intermediate features, or a dictionary containing
+            'image_features' and 'image_intermediates' (and optionally 'image_intermediates_prefix')
+
+        """
+        assert output_fmt in ('NCHWD', 'NLC'), 'Output format must be one of NCHWD or NLC.'
+        reshape = output_fmt == 'NCHWD'
+        intermediates = []
+        take_indices, max_index = feature_take_indices(len(self.blocks), indices)
+
+        # forward pass
+        B, _, height, width, depth = x.shape
+        x = self.patch_embed(x)
+        x, rope = self._pos_embed(x)
+        x = self.patch_drop(x)
+        x = self.norm_pre(x)
+
+        if torch.jit.is_scripting() or not stop_early:  # can't slice blocks in torchscript
+            blocks = self.blocks
+        else:
+            blocks = self.blocks[:max_index + 1]
+        for i, blk in enumerate(blocks):
+            x = blk(x, rope=rope)
+            if i in take_indices:
+                # normalize intermediates with final norm layer if enabled
+                intermediates.append(self.norm(x) if norm else x)
+
+        # process intermediates
+        if self.num_prefix_tokens:
+            # split prefix (e.g. class, distill) and spatial feature tokens
+            prefix_tokens = [y[:, 0:self.num_prefix_tokens] for y in intermediates]
+            intermediates = [y[:, self.num_prefix_tokens:] for y in intermediates]
+        else:
+            prefix_tokens = None
+
+        if reshape:
+            # reshape to BCHW output format
+            H, W, D = self.patch_embed.dynamic_feat_size((height, width, depth))
+            intermediates = [y.reshape(B, H, W, D, -1).permute(0, 4, 1, 2, 3).contiguous() for y in intermediates]
+
+        if output_dict:
+            result_dict = {}
+            # Intermediates are always included
+            result_dict['image_intermediates'] = intermediates
+            if prefix_tokens is not None and return_prefix_tokens:
+                result_dict['image_intermediates_prefix'] = prefix_tokens
+            
+            # Only include features if not intermediates_only
+            if not intermediates_only:
+                x_final = self.norm(x)
+                result_dict['image_features'] = x_final
+
+            return result_dict
+
+        # For non-dictionary output, maintain the original behavior
+        if not torch.jit.is_scripting() and return_prefix_tokens and prefix_tokens is not None:
+            # return_prefix not support in torchscript due to poor type handling
+            intermediates = list(zip(intermediates, prefix_tokens))
+
+        if intermediates_only:
+            return intermediates
+
+        x = self.norm(x)
+
+        return x, intermediates
+
+    def prune_intermediate_layers(
+        self,
+        indices: Union[int, List[int]] = 1,
+        prune_norm: bool = False,
+        prune_head: bool = True,
+    ):
+        """Prune layers not required for specified intermediates.
+
+        Args:
+            indices: Indices of intermediate layers to keep.
+            prune_norm: Whether to prune normalization layer.
+            prune_head: Whether to prune the classifier head.
+
+        Returns:
+            List of indices that were kept.
+        """
+        take_indices, max_index = feature_take_indices(len(self.blocks), indices)
+        self.blocks = self.blocks[:max_index + 1]  # truncate blocks
+        if prune_norm:
+            self.norm = nn.Identity()
+        if prune_head:
+            self.fc_norm = nn.Identity()
+            self.reset_classifier(0, '')
+        return take_indices
+    
+    def get_intermediate_layers(
+        self,
+        x: torch.Tensor,
+        n: Union[int, List[int], Tuple[int]] = 1,
+        reshape: bool = False,
+        return_prefix_tokens: bool = False,
+        norm: bool = False,
+    ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]:
+        """Get intermediate layer outputs (DINO interface compatibility).
+
+        NOTE: This API is for backwards compat, favour using forward_intermediates() directly.
+
+        Args:
+            x: Input tensor.
+            n: Number or indices of layers.
+            reshape: Reshape to NCHWD format.
+            return_prefix_tokens: Return prefix tokens.
+            norm: Apply normalization.
+
+        Returns:
+            List of intermediate features.
+        """
+        return self.forward_intermediates(
+            x, n,
+            return_prefix_tokens=return_prefix_tokens,
+            norm=norm,
+            output_fmt='NCHWD' if reshape else 'NLC',
+            intermediates_only=True,
+        )
+    
+    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass through feature layers (embeddings, transformer blocks, post-transformer norm)."""
+        x = self.patch_embed(x)
+        x, rope = self._pos_embed(x)
+        x = self.patch_drop(x)
+        x = self.norm_pre(x)
+
+        for blk in self.blocks:
+            x = blk(x, rope=rope)
+        x = self.norm(x)
+        return x
+
+    def pool(self, x: torch.Tensor, pool_type: Optional[str] = None) -> torch.Tensor:
+        """Apply pooling to feature tokens.
+
+        Args:
+            x: Feature tensor.
+            pool_type: Pooling type override.
+
+        Returns:
+            Pooled features.
+        """
+        if self.attn_pool is not None:
+            x = self.attn_pool(x)
+            return x
+        pool_type = self.global_pool if pool_type is None else pool_type
+        x = global_pool_nlc(
+            x, 
+            pool_type=pool_type, 
+            num_prefix_tokens=self.num_prefix_tokens,
+        )
+        return x
+
+    def forward_head(self, x: torch.Tensor, pre_logits: bool = False) -> torch.Tensor:
+        """Forward pass through classifier head.
+
+        Args:
+            x: Feature tensor.
+            pre_logits: Return features before final classifier.
+
+        Returns:
+            Output tensor.
+        """
+        x = self.pool(x)
+        x = self.fc_norm(x)
+        x = self.head_drop(x)
+        return x if pre_logits else self.head(x)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.forward_features(x)
+        x = self.forward_head(x)
+        return x
+    
+    @classmethod
+    def from_pretrained(
+            cls,
+            checkpoint_path_or_url: Union[str, os.PathLike],
+            verbose: 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
+
+
+def vit_tiny_patch16_128(
+    checkpoint_path_or_url: Optional[str] = None, 
+    **kwargs
+) -> VisionTransformer:
+    """ViT-Tiny model with 3D patch embedding, patch size [16, 16, 8] and input size [128, 128, 64].
+    """
+    kwargs = dict(
+        img_size=(128, 128, 64),
+        patch_size=(16, 16, 8),
+        embed_dim=192,
+        depth=12,
+        num_heads=2,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=nn.LayerNorm,
+        **kwargs,
+    )
+    if checkpoint_path_or_url is not None:
+        return VisionTransformer.from_pretrained(checkpoint_path_or_url, **kwargs)
+    return VisionTransformer(**kwargs)
+
+
+def vit_small_patch16_128(
+    checkpoint_path_or_url: Optional[str] = None, 
+    **kwargs
+) -> VisionTransformer:
+    """ViT-Small model with 3D patch embedding, patch size [16, 16, 8] and input size [128, 128, 64].
+    """
+    kwargs = dict(
+        img_size=(128, 128, 64),
+        patch_size=(16, 16, 8),
+        embed_dim=384,
+        depth=12,
+        num_heads=4,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=nn.LayerNorm,
+        **kwargs,
+    )
+    if checkpoint_path_or_url is not None:
+        return VisionTransformer.from_pretrained(checkpoint_path_or_url, **kwargs)
+    return VisionTransformer(**kwargs)
+
+
+def vit_base_patch16_128(
+    checkpoint_path_or_url: Optional[str] = None, 
+    **kwargs
+) -> VisionTransformer:
+    """ViT-Base model with 3D patch embedding, patch size [16, 16, 8] and input size [128, 128, 64].
+    """
+    kwargs = dict(
+        img_size=(128, 128, 64),
+        patch_size=(16, 16, 8),
+        embed_dim=768,
+        depth=12,
+        num_heads=8,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=nn.LayerNorm,
+        **kwargs,
+    )
+    if checkpoint_path_or_url is not None:
+        return VisionTransformer.from_pretrained(checkpoint_path_or_url, **kwargs)
+    return VisionTransformer(**kwargs)
+
+def vit_base_patch16_256(
+    pretrained_weights: Optional[str] = None, 
+    **kwargs
+) -> VisionTransformer:
+    """ViT-Base model with 3D patch embedding, patch size [16, 16, 8] and input size [256, 256, 128].
+    """
+    kwargs = dict(
+        img_size=(256, 256, 128),
+        patch_size=(16, 16, 8),
+        embed_dim=768,
+        depth=12,
+        num_heads=8,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=nn.LayerNorm,
+        **kwargs,
+    )
+    if pretrained_weights is not None:
+        return VisionTransformer.from_pretrained(pretrained_weights, **kwargs)
+    return VisionTransformer(**kwargs)
+    
+
+def vit_base_patch32_128(
+    checkpoint_path_or_url: Optional[str] = None, 
+    **kwargs
+) -> VisionTransformer:
+    """ViT-Base model with 3D patch embedding, patch size [32, 32, 16] and input size [128, 128, 64].
+    """
+    kwargs = dict(
+        img_size=(128, 128, 64),
+        patch_size=(32, 32, 16),
+        embed_dim=768,
+        depth=12,
+        num_heads=8,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=nn.LayerNorm,
+        **kwargs,
+    )
+    if checkpoint_path_or_url is not None:
+        return VisionTransformer.from_pretrained(checkpoint_path_or_url, **kwargs)
+    return VisionTransformer(**kwargs)
+
+
+def vit_large_patch16_128(
+    checkpoint_path_or_url: Optional[str] = None, 
+    **kwargs
+) -> VisionTransformer:
+    """ViT-Large model with 3D patch embedding, patch size [16, 16, 8] and input size [128, 128, 64].
+    """
+    kwargs = dict(
+        img_size=(128, 128, 64),
+        patch_size=(16, 16, 8),
+        embed_dim=1080,
+        depth=24,
+        num_heads=12,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=nn.LayerNorm,
+        **kwargs,
+    )
+    if checkpoint_path_or_url is not None:
+        return VisionTransformer.from_pretrained(checkpoint_path_or_url, **kwargs)
+    return VisionTransformer(**kwargs)
+
+def vit_large_patch16_256(
+    pretrained_weights: Optional[str] = None, 
+    **kwargs
+) -> VisionTransformer:
+    """ViT-Large model with 3D patch embedding, patch size [16, 16, 8] and input size [128, 128, 64].
+    """
+    kwargs = dict(
+        img_size=(256, 256, 128),
+        patch_size=(16, 16, 8),
+        embed_dim=1080,
+        depth=24,
+        num_heads=12,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=nn.LayerNorm,
+        **kwargs,
+    )
+    if pretrained_weights is not None:
+        return VisionTransformer.from_pretrained(pretrained_weights, **kwargs)
+    return VisionTransformer(**kwargs)
+
+def vit_large_patch16_320(
+    pretrained_weights: Optional[str] = None, 
+    **kwargs
+) -> VisionTransformer:
+    """ViT-Large model with 3D patch embedding, patch size [16, 16, 8] and input size [320, 320, 128].
+    """
+    kwargs = dict(
+        img_size=(320, 320, 128),
+        patch_size=(16, 16, 8),
+        embed_dim=1080,
+        depth=24,
+        num_heads=12,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=nn.LayerNorm,
+        **kwargs,
+    )
+    if pretrained_weights is not None:
+        return VisionTransformer.from_pretrained(pretrained_weights, **kwargs)
+    return VisionTransformer(**kwargs)
+
+
+def vit_large_patch32_128(
+    checkpoint_path_or_url: Optional[str] = None, 
+    **kwargs
+) -> VisionTransformer:
+    """ViT-Large model with 3D patch embedding, patch size [32, 32, 16] and input size [128, 128, 64].
+    """
+    kwargs = dict(
+        img_size=(128, 128, 64),
+        patch_size=(32, 32, 16),
+        embed_dim=1080,
+        depth=24,
+        num_heads=12,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=nn.LayerNorm,
+        **kwargs,
+    )
+    if checkpoint_path_or_url is not None:
+        return VisionTransformer.from_pretrained(checkpoint_path_or_url, **kwargs)
+    return VisionTransformer(**kwargs)

spectre/models/vision_transformer_features.py

@@ -0,0 +1,455 @@
+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.models.vision_transformer import Block
+from spectre.models.layers import RotaryPositionEmbedding
+
+
+class FeatureVisionTransformer(nn.Module):
+    """ Vision Transformer that accepts flattened patches as input.
+    """
+    def __init__(
+        self, 
+        grid_size: Optional[Union[int, Tuple[int, int, int]]] = None,
+        patch_dim: int = 768,
+        num_classes: int = 1000,
+        global_pool: Literal['', 'avg', 'avgmax', 'max', 'token', 'map'] = 'token',
+        embed_dim: int = 768,
+        depth: int = 12,
+        num_heads: int = 12,
+        attn_mode: str = 'mha',
+        q_proj_dim: Optional[int] = None,
+        kv_proj_dim: Optional[int] = None,
+        mlp_ratio: float = 4.,
+        qkv_bias: bool = True,
+        qk_norm: bool = False,
+        proj_bias: bool = True,
+        init_values: Optional[float] = None,
+        class_token: bool = True,
+        pos_embed: str = 'learn',
+        no_embed_class: bool = False,
+        rope_kwargs: Optional[dict] = None,
+        reg_tokens: int = 0,
+        pre_norm: bool = False,
+        final_norm: bool = True,
+        fc_norm: Optional[bool] = None,
+        dynamic_grid_size: bool = False,
+        drop_rate: float = 0.,
+        pos_drop_rate: float = 0.,
+        patch_drop_rate: float = 0.,
+        proj_drop_rate: float = 0.,
+        attn_drop_rate: float = 0.,
+        drop_path_rate: float = 0.,
+        norm_layer: Optional[Union[Callable, Type[torch.nn.Module]]] = None,
+        act_layer: Optional[Union[Callable, Type[torch.nn.Module]]] = None,
+        block_fn: Type[nn.Module] = Block,
+        mlp_layer: Type[nn.Module] = Mlp,
+    ) -> None:
+        """
+        Args:
+            num_patches: Number of patches in the input.
+            patch_dim: Dimension of each flattened input patch.
+            num_classes: Number of classes for classification head.
+            global_pool: Type of global pooling for final sequence (default: 'token').
+            embed_dim: Transformer embedding dimension.
+            depth: Depth of transformer.
+            num_heads: Number of attention heads.
+            attn_mode: Attention mode ('mha', 'mqa', 'mla').
+            q_proj_dim: Query projection dimension for 'mla' mode.
+            kv_proj_dim: Key, value projection dimension for 'mla' mode.
+            mlp_ratio: Ratio of mlp hidden dim to embedding dim.
+            qkv_bias: Enable bias for qkv projections if True.
+            init_values: Layer-scale init values (layer-scale enabled if not None).
+            class_token: Use class token.
+            no_embed_class: Don't include position embeddings for class (or reg) tokens.
+            reg_tokens: Number of register tokens.
+            pre_norm: Enable norm after embeddings, before transformer blocks (standard in CLIP ViT).
+            final_norm: Enable norm after transformer blocks, before head (standard in most ViT).
+            fc_norm: Move final norm after pool (instead of before), if None, enabled when global_pool == 'avg'.
+            drop_rate: Head dropout rate.
+            pos_drop_rate: Position embedding dropout rate.
+            attn_drop_rate: Attention dropout rate.
+            drop_path_rate: Stochastic depth rate.
+            weight_init: Weight initialization scheme.
+            fix_init: Apply weight initialization fix (scaling w/ layer index).
+            norm_layer: Normalization layer.
+            act_layer: MLP activation layer.
+            block_fn: Transformer block layer.
+        """
+        super().__init__()
+        assert global_pool in ('', 'avg', 'avgmax', 'max', 'token', 'map')
+        assert class_token or global_pool != 'token'
+        assert pos_embed in ('', 'none', 'learn', 'rope')
+        assert attn_mode in ('mha', 'mqa', 'mla')
+        assert grid_size is not None or pos_embed in ('', 'none', 'rope')
+        rope_kwargs = {} if rope_kwargs is None else dict(rope_kwargs)
+        rope_kwargs.setdefault("dtype", torch.float32)  # robust with mixed-precision
+        use_fc_norm = global_pool in ('avg', 'avgmax', 'max') if fc_norm is None else fc_norm
+        norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
+        act_layer = act_layer or nn.GELU
+
+        self.grid_size = None if grid_size is None else to_3tuple(grid_size)
+        self.num_classes = num_classes
+        self.global_pool = global_pool
+        self.num_features = self.head_hidden_size = self.embed_dim = embed_dim  # for consistency with other models
+        self.num_prefix_tokens = 1 if class_token else 0
+        self.num_prefix_tokens += reg_tokens
+        self.num_reg_tokens = reg_tokens
+        self.has_class_token = class_token
+        self.no_embed_class = no_embed_class  # don't embed prefix positions (includes reg)
+        self.dynamic_grid_size = dynamic_grid_size
+
+        self.num_patches = None if grid_size is None else int(math.prod(grid_size))
+        self.patch_proj = nn.Linear(patch_dim, embed_dim, proj_bias)
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) if class_token else None
+        self.reg_token = nn.Parameter(torch.zeros(1, reg_tokens, embed_dim)) if reg_tokens else None
+        self.pos_embed, self.rope, self.requires_per_sample_rope = None, None, False
+        if pos_embed == 'learn':
+            embed_len = self.num_patches if no_embed_class else self.num_patches + self.num_prefix_tokens
+            self.pos_embed = nn.Parameter(torch.randn(1, embed_len, embed_dim) * .02)
+        if pos_embed == 'rope':
+            self.rope = RotaryPositionEmbedding(
+                embed_dim=embed_dim,
+                num_heads=num_heads,
+                **rope_kwargs,
+            )
+        self.pos_drop = nn.Dropout(p=pos_drop_rate)
+        if patch_drop_rate > 0:
+            self.patch_drop = PatchDropout(
+                patch_drop_rate,
+                num_prefix_tokens=self.num_prefix_tokens,
+            )
+        else:
+            self.patch_drop = nn.Identity()
+        self.norm_pre = norm_layer(embed_dim) if pre_norm else nn.Identity()
+
+        dpr = [drop_path_rate * i / (depth - 1) if depth > 1 else 0.0 for i in range(depth)]  # stochastic depth decay rule
+        self.blocks = nn.Sequential(*[
+            block_fn(
+                dim=embed_dim,
+                num_heads=num_heads,
+                attn_mode=attn_mode,
+                q_proj_dim=q_proj_dim,
+                kv_proj_dim=kv_proj_dim,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                qk_norm=qk_norm,
+                proj_bias=proj_bias,
+                init_values=init_values,
+                proj_drop=proj_drop_rate,
+                attn_drop=attn_drop_rate,
+                drop_path=dpr[i],
+                norm_layer=norm_layer,
+                act_layer=act_layer,
+                mlp_layer=mlp_layer,
+            )
+            for i in range(depth)])
+        self.feature_info = [
+            dict(module=f'blocks.{i}', num_chs=embed_dim) for i in range(depth)]
+        self.norm = norm_layer(embed_dim) if final_norm and not use_fc_norm else nn.Identity()
+
+        # Classifier Head
+        if global_pool == 'map':
+            self.attn_pool = AttentionPoolLatent(
+                self.embed_dim,
+                num_heads=num_heads,
+                mlp_ratio=mlp_ratio,
+                norm_layer=norm_layer,
+                act_layer=act_layer,
+            )
+        else:
+            self.attn_pool = None
+        self.fc_norm = norm_layer(embed_dim) if final_norm and use_fc_norm else nn.Identity()
+        self.head_drop = nn.Dropout(drop_rate)
+        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+        self.init_weights()
+
+    def init_weights(self) -> None:
+        if self.pos_embed is not None and not self.pos_embed.is_meta:
+            nn.init.trunc_normal_(self.pos_embed, std=.02)
+        if self.cls_token is not None and not self.cls_token.is_meta:
+            nn.init.normal_(self.cls_token, std=1e-6)
+        if self.reg_token is not None and not self.reg_token.is_meta:
+            nn.init.normal_(self.reg_token, std=1e-6)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m: nn.Module) -> None:
+        # this fn left here for compat with downstream users
+        if isinstance(m, nn.Linear):
+            if not m.weight.is_meta:
+                nn.init.trunc_normal_(m.weight, std=.02)
+            if m.bias is not None and not m.bias.is_meta:
+                nn.init.zeros_(m.bias)
+
+    @torch.jit.ignore
+    def no_weight_decay(self) -> Set:
+        return {'pos_embed', 'cls_token', 'dist_token'}
+
+    @torch.jit.ignore
+    def get_classifier(self) -> nn.Module:
+        return self.head
+
+    def reset_classifier(self, num_classes: int, global_pool: Optional[str] = None):
+        self.num_classes = num_classes
+        if global_pool is not None:
+            assert global_pool in ('', 'avg', 'avgmax', 'max', 'token', 'map')
+            if global_pool == 'map' and self.attn_pool is None:
+                assert False, "Cannot currently add attention pooling in reset_classifier()."
+            elif global_pool != 'map' and self.attn_pool is not None:
+                self.attn_pool = None  # remove attention pooling
+            self.global_pool = global_pool
+        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+    def _pos_embed(
+        self, 
+        x: torch.Tensor,
+        grid_size: Optional[Union[int, Tuple[int, int, int]]] = None,
+    ):
+        if self.pos_embed is None and self.rope is None:
+            x = x.view(x.shape[0], -1, x.shape[-1])
+            if self.reg_token is not None:
+                x = torch.cat([self.reg_token.expand(x.shape[0], -1, -1), x], dim=1)
+            if self.cls_token is not None:
+                x = torch.cat([self.cls_token.expand(x.shape[0], -1, -1), x], dim=1)
+            return x, None
+        
+        if self.dynamic_grid_size or self.rope is not None:
+            assert grid_size is not None, "grid_size must be provided when using dynamic_grid_size or RoPE."
+        
+        pos_embed, rope = None, None
+        if self.pos_embed is not None:
+            if self.dynamic_grid_size:
+                H, W, D = to_3tuple(grid_size)
+                prev_grid_size = self.grid_size
+                pos_embed = resample_abs_pos_embed(
+                    self.pos_embed, 
+                    new_size=(H, W, D),
+                    old_size=prev_grid_size,
+                    num_prefix_tokens=0 if self.no_embed_class else self.num_prefix_tokens,
+                )
+            else:
+                pos_embed = self.pos_embed
+
+        if self.rope is not None:
+            B = x.shape[0]            
+            H, W, D = to_3tuple(grid_size)
+            if self.requires_per_sample_rope:
+                rope = [self.rope(H=H, W=W, D=D) for _ in range(B)]
+            else:
+                rope = self.rope(H=H, W=W, D=D)
+
+        to_cat = []
+        if self.cls_token is not None:
+            to_cat.append(self.cls_token.expand(x.shape[0], -1, -1))
+        if self.reg_token is not None:
+            to_cat.append(self.reg_token.expand(x.shape[0], -1, -1))
+
+        if self.no_embed_class:
+            # deit-3, updated JAX (big vision)
+            # position embedding does not overlap with class token, add then concat
+            if pos_embed is not None:
+                x = x + pos_embed
+            if to_cat:
+                x = torch.cat(to_cat + [x], dim=1)
+        else:
+            # original timm, JAX, and deit vit impl
+            # pos_embed has entry for class token, concat then add
+            if to_cat:
+                x = torch.cat(to_cat + [x], dim=1)
+            if pos_embed is not None:
+                x = x + pos_embed
+
+        return self.pos_drop(x), rope
+
+    def forward_features(
+        self, 
+        x: torch.Tensor, 
+        grid_size: Optional[Union[int, Tuple[int, int, int]]] = None,
+    ) -> torch.Tensor:
+        assert x.ndim == 3, f"Expected input with 3 dimensions (B, N, C), got {x.ndim}."
+
+        x = self.patch_proj(x)
+        x, rope = self._pos_embed(x, grid_size)
+        x = self.patch_drop(x)
+        x = self.norm_pre(x)
+
+        for blk in self.blocks:
+            x = blk(x, rope=rope)
+        x = self.norm(x)
+        return x
+
+    def pool(self, x: torch.Tensor, pool_type: Optional[str] = None) -> torch.Tensor:
+        if self.attn_pool is not None:
+            x = self.attn_pool(x)
+            return x
+        pool_type = self.global_pool if pool_type is None else pool_type
+        x = global_pool_nlc(x, pool_type=pool_type, num_prefix_tokens=self.num_prefix_tokens)
+        return x
+
+    def forward_head(self, x: torch.Tensor, pre_logits: bool = False) -> torch.Tensor:
+        x = self.pool(x)
+        x = self.fc_norm(x)
+        x = self.head_drop(x)
+        return x if pre_logits else self.head(x)
+
+    def forward(
+        self, 
+        x: torch.Tensor, 
+        grid_size: Optional[Union[int, Tuple[int, int, int]]] = None,
+    ) -> torch.Tensor:
+        x = self.forward_features(x, grid_size)
+        x = self.forward_head(x)
+        return x
+    
+    @classmethod
+    def from_pretrained(
+            cls,
+            checkpoint_path_or_url: Union[str, os.PathLike],
+            verbose: 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
+
+
+def feat_vit_tiny(
+    patch_dim,
+    checkpoint_path_or_url: Optional[str] = None,
+    **kwargs,
+) -> FeatureVisionTransformer:
+    """Feature ViT-Tiny model.
+    """
+    kwargs = dict(
+        patch_dim=patch_dim,
+        embed_dim=192,
+        depth=2,
+        num_heads=2,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=nn.LayerNorm,
+        **kwargs,
+    )
+    if checkpoint_path_or_url is not None:
+        return FeatureVisionTransformer.from_pretrained(checkpoint_path_or_url, **kwargs)
+    return FeatureVisionTransformer(**kwargs)
+
+
+def feat_vit_small(
+    patch_dim,
+    checkpoint_path_or_url: Optional[str] = None,
+    **kwargs,
+) -> FeatureVisionTransformer:
+    """Feature ViT-Small model.
+    """
+    kwargs = dict(
+        patch_dim=patch_dim,
+        embed_dim=384,
+        depth=2,
+        num_heads=4,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=nn.LayerNorm,
+        **kwargs,
+    )
+    if checkpoint_path_or_url is not None:
+        return FeatureVisionTransformer.from_pretrained(checkpoint_path_or_url, **kwargs)
+    return FeatureVisionTransformer(**kwargs)
+
+
+def feat_vit_base(
+    patch_dim,
+    checkpoint_path_or_url: Optional[str] = None,
+    **kwargs,
+) -> FeatureVisionTransformer:
+    """Feature ViT-Base model.
+    """
+    kwargs = dict(
+        patch_dim=patch_dim,
+        embed_dim=768,
+        depth=2,
+        num_heads=8,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=nn.LayerNorm,
+        **kwargs,
+    )
+    if checkpoint_path_or_url is not None:
+        return FeatureVisionTransformer.from_pretrained(checkpoint_path_or_url, **kwargs)
+    return FeatureVisionTransformer(**kwargs)
+
+
+def feat_vit_large(
+    patch_dim,
+    checkpoint_path_or_url: Optional[str] = None,
+    **kwargs,
+) -> FeatureVisionTransformer:
+    """Feature ViT-Large model.
+    """
+    kwargs = dict(
+        patch_dim=patch_dim,
+        embed_dim=1080,
+        depth=4,
+        num_heads=12,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=nn.LayerNorm,
+        **kwargs,
+    )
+    if checkpoint_path_or_url is not None:
+        return FeatureVisionTransformer.from_pretrained(checkpoint_path_or_url, **kwargs)
+    return FeatureVisionTransformer(**kwargs)

spectre/utils/__init__.py

@@ -0,0 +1,117 @@
+from ._utils import (
+    fix_random_seeds, 
+    to_ntuple,
+    to_1tuple,
+    to_2tuple,
+    to_3tuple,
+    to_4tuple,
+)
+from .checkpointing import (
+    save_state,
+    load_state,
+    extract_model_from_checkpoint_dinov2,
+    extract_model_from_checkpoint_siglip,
+)
+from .collate import (
+    extended_collate_dino,
+    extended_collate_siglip,
+    collate_add_filenames,
+)
+from .config import setup
+from .dataloader import get_dataloader
+from .distributed import (
+    is_enabled,
+    get_global_size,
+    get_global_rank,
+    get_local_size,
+    get_local_rank,
+    init_distributed,
+)
+from .lora import add_lora_adapters
+from .masking import random_block_mask
+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,
+)
+from .param_groups import get_param_groups_with_decay
+from .scheduler import (
+    linear_warmup_schedule,
+    cosine_schedule,
+    cosine_warmup_schedule,
+    CosineWarmupScheduler,
+)
+
+__all__ = [
+    "fix_random_seeds",
+    "to_ntuple",
+    "to_1tuple",
+    "to_2tuple",
+    "to_3tuple",
+    "to_4tuple",
+    "save_state",
+    "load_state",
+    "extract_model_from_checkpoint_dinov2",
+    "extract_model_from_checkpoint_siglip",
+    "extended_collate_dino",
+    "extended_collate_siglip",
+    "collate_add_filenames",
+    "setup",
+    "get_dataloader",
+    "is_enabled",
+    "get_global_size",
+    "get_global_rank",
+    "get_local_size",
+    "get_local_rank",
+    "init_distributed",
+    "add_lora_adapters",
+    "random_block_mask",
+    "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",
+    "get_param_groups_with_decay",
+    "linear_warmup_schedule",
+    "cosine_schedule",
+    "cosine_warmup_schedule",
+    "CosineWarmupScheduler",
+]

spectre/utils/_utils.py

@@ -0,0 +1,49 @@
+import random
+from typing import Iterable
+from itertools import repeat
+
+import torch
+import numpy as np
+
+MONAI_IMPORT_ERROR = None
+try:
+    import monai
+except ImportError as e:
+    monai = None  # type: ignore
+    MONAI_IMPORT_ERROR = e
+
+def fix_random_seeds(seed: int = 31):
+    """ 
+    Fix random seeds.
+    """
+    if MONAI_IMPORT_ERROR is not None:
+        raise ImportError(
+            "MONAI is required to use fix_random_seeds but not installed. "
+            "Please install MONAI to use this function."
+        ) from MONAI_IMPORT_ERROR
+
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+    monai.utils.set_determinism(seed=seed)
+
+
+def _ntuple(n: int):
+    """
+    Helper function to create n-tuple.
+    """
+    def parse(x):
+        if isinstance(x, Iterable) and not isinstance(x, str):
+            return tuple(x)
+        return tuple(repeat(x, n))
+    return parse
+
+
+to_1tuple = _ntuple(1)
+to_2tuple = _ntuple(2)
+to_3tuple = _ntuple(3)
+to_4tuple = _ntuple(4)
+to_ntuple = _ntuple

spectre/utils/checkpointing.py

@@ -0,0 +1,238 @@
+import os
+import random
+import warnings
+from typing import Optional, Any
+
+import torch
+import numpy as np
+import torch.distributed as dist
+
+
+def _get_local_rng_state() -> dict:
+    """Return a picklable dict with local RNG states (cpu & cuda, numpy, random)."""
+    state = {
+        "torch": torch.get_rng_state().cpu(),
+        "numpy": np.random.get_state(),
+        "random": random.getstate(),
+    }
+
+    if torch.cuda.is_available():
+        # make sure CUDA states are stored on CPU so they are picklable
+        cuda_states = [s.cpu() for s in torch.cuda.get_rng_state_all()]
+        state["cuda"] = cuda_states
+    else:
+        state["cuda"] = None
+
+    return state
+
+
+def _set_local_rng_state(state: dict) -> None:
+    """Set local RNG states from the dict produced by _get_local_rng_state()."""
+    if state is None:
+        return
+
+    if "torch" in state and state["torch"] is not None:
+        torch.set_rng_state(state["torch"])
+    if "cuda" in state and state["cuda"] is not None and torch.cuda.is_available():
+        try:
+            # move back to CUDA tensors for this process and set them
+            cuda_states = [s.cuda() for s in state["cuda"]]
+            torch.cuda.set_rng_state_all(cuda_states)
+        except Exception:
+            # fallback: try setting per-device RNG if set_rng_state_all fails
+            for i, s in enumerate(state["cuda"]):
+                try:
+                    torch.cuda.set_rng_state(s.cuda(), device=i)
+                except Exception:
+                    # ignore if device mismatch
+                    pass
+
+    if "numpy" in state and state["numpy"] is not None:
+        np.random.set_state(state["numpy"])
+    if "random" in state and state["random"] is not None:
+        random.setstate(state["random"])
+
+
+def save_state(ckpt_path: str, epoch: Optional[int] = None, **named_objects: Any) -> None:
+    """
+    Save a checkpoint that includes:
+      - epoch (optional)
+      - state_dicts for provided named_objects
+      - rng_states: list of per-rank RNG dictionaries (one entry per world rank)
+
+    If torch.distributed is initialized the RNG states from all ranks are gathered and
+    stored in checkpoint["rng_states"] (list indexed by rank). Only rank 0 writes the file.
+    In single-process mode the checkpoint contains a single-item rng_states list.
+    """
+    os.makedirs(os.path.dirname(ckpt_path), exist_ok=True)
+
+    # prepare local RNG state
+    local_rng = _get_local_rng_state()
+
+    # distributed path: gather RNG states from all ranks
+    if dist.is_available() and dist.is_initialized():
+        rank = dist.get_rank()
+        world_size = dist.get_world_size()
+        all_states = [None] * world_size
+        # gather python objects (picklable)
+        dist.all_gather_object(all_states, local_rng)
+
+        # only rank 0 writes the checkpoint file
+        if rank == 0:
+            checkpoint = {}
+            if epoch is not None:
+                checkpoint["epoch"] = epoch
+            checkpoint["rng_states"] = all_states
+
+            # save provided objects' state_dicts (rank 0's local state_dicts)
+            for name, obj in named_objects.items():
+                checkpoint[name] = obj.state_dict()
+
+            torch.save(checkpoint, ckpt_path)
+
+        # ensure everyone waits until rank 0 finished writing
+        dist.barrier()
+
+    else:
+        # single-process fallback
+        checkpoint = {}
+        if epoch is not None:
+            checkpoint["epoch"] = epoch
+        checkpoint["rng_states"] = [local_rng]
+        for name, obj in named_objects.items():
+            checkpoint[name] = obj.state_dict()
+        torch.save(checkpoint, ckpt_path)
+
+
+def load_state(ckpt_path: str, **named_objects: Any) -> int:
+    """
+    Load checkpoint saved by save_state.
+
+    - Each process loads the same file and restores its own RNG state (checkpoint['rng_states'][rank]).
+    - Named objects that exist in the checkpoint will have their state_dict loaded.
+    - Returns epoch (int) if present, otherwise 0.
+    """
+    if not os.path.isfile(ckpt_path):
+        warnings.warn(f"Checkpoint file not found: {ckpt_path}")
+        return 0
+
+    # load on all ranks (shared FS assumed)
+    checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+    epoch = checkpoint.get("epoch", 0)
+
+    # load state_dicts into provided objects
+    for name, obj in named_objects.items():
+        if name in checkpoint:
+            try:
+                obj.load_state_dict(checkpoint[name])
+            except Exception as e:
+                warnings.warn(f"Failed to load state_dict for '{name}': {e}")
+        else:
+            warnings.warn(f"No state_dict found for '{name}' in checkpoint.")
+
+    # restore this rank's RNG state
+    rng_states = checkpoint.get("rng_states", None)
+    if rng_states is not None:
+        if dist.is_available() and dist.is_initialized():
+            rank = dist.get_rank()
+            if rank < len(rng_states):
+                my_state = rng_states[rank]
+            else:
+                my_state = None
+        else:
+            # single-process file: first element
+            my_state = rng_states[0] if len(rng_states) > 0 else None
+
+        try:
+            _set_local_rng_state(my_state)
+        except Exception as e:
+            warnings.warn(f"Failed to restore RNG state: {e}")
+
+    else:
+        warnings.warn("No 'rng_states' found in checkpoint; RNGs not restored.")
+
+    return epoch
+
+
+def extract_model_from_checkpoint_dinov2(checkpoint_path: str):
+    # Load the checkpoint
+    checkpoint = torch.load(
+        checkpoint_path, 
+        weights_only=False, 
+        map_location="cpu"
+    )
+    
+    # Get model state dict
+    model_state = checkpoint.get("model", checkpoint)
+
+    # Create output folder
+    output_dir = str(checkpoint_path).replace(".pt", "")
+    os.makedirs(output_dir, exist_ok=True)
+
+    # Quick check: compare the parameters of head_teacher_ibot vs head_teacher_dino
+    teacher_dino_keys = [k for k in model_state.keys() if k.startswith("head_teacher_dino.")]
+    teacher_ibot_keys = [k for k in model_state.keys() if k.startswith("head_teacher_ibot.")]
+
+    ibot_separate = True
+    if teacher_dino_keys and teacher_ibot_keys:
+        if all(torch.equal(model_state[dino_key], model_state[ibot_key]) \
+               for dino_key, ibot_key in zip(teacher_dino_keys, teacher_ibot_keys)):
+            ibot_separate = False  # Same weights → no separate ibot head
+    
+    # Define the components to save
+    components = {
+        "backbone_teacher.pt": "backbone_teacher.vit",
+        "backbone_student.pt": "backbone_student.vit",
+        "head_student_dino.pt": "head_student_dino",
+        "head_teacher_dino.pt": "head_teacher_dino"
+    }
+
+    # Add ibot heads only if separate
+    if ibot_separate:
+        components["head_student_ibot.pt"] = "head_student_ibot"
+        components["head_teacher_ibot.pt"] = "head_teacher_ibot"
+
+    # Extract and save each component
+    for filename, key in components.items():
+        sub_state_dict = {k.replace(f"{key}.", ""): v for k, v in model_state.items() if k.startswith(key)}
+        if not sub_state_dict:
+            print(f"[WARNING] No parameters found for {key}, skipping...")
+            continue
+        torch.save(sub_state_dict, os.path.join(output_dir, filename))
+    
+    print(f"Components extracted to: {output_dir}")
+
+
+def extract_model_from_checkpoint_siglip(checkpoint_path: str):
+    # Load the checkpoint
+    checkpoint = torch.load(
+        checkpoint_path, 
+        weights_only=False, 
+        map_location="cpu",
+    )
+    
+    # Get model state dict
+    model_state = checkpoint.get("model", checkpoint)
+
+    # Create output folder
+    output_dir = str(checkpoint_path).replace(".pt", "")
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # Define the components to save
+    components = {
+        "backbone_image.pt": "backbone_image",
+        "backbone_text.pt": "backbone_text",
+        "feature_comb_image.pt": "feature_comb_image",
+        "projection_image.pt": "projection_image",
+        "projection_text.pt": "projection_text"
+    }
+
+    # Extract and save each component
+    for filename, key in components.items():
+        sub_state_dict = {k.replace(f"{key}.", ""): v for k, v in model_state.items() if k.startswith(key)}
+        if not sub_state_dict:
+            print(f"[WARNING] No parameters found for {key}, skipping...")
+            continue
+        torch.save(sub_state_dict, os.path.join(output_dir, filename))
+    
+    print(f"Components extracted to: {output_dir}")

spectre/utils/collate.py

@@ -0,0 +1,120 @@
+from typing import List, Callable, Optional
+
+import torch
+
+MONAI_IMPORT_ERROR = None
+try:
+    from monai.data import list_data_collate
+except ImportError as e:
+    list_data_collate = lambda x: x  # type: ignore
+    MONAI_IMPORT_ERROR = e
+
+
+def extended_collate_dino(samples_list: List) -> dict:
+    """
+    Applies MONAI's list_data_collate first and then extends it with DINOv2 masking logic.
+
+    Args:
+        samples_list: List of samples containing 'global_crops' and 'local_crops'.
+        mask_ratio: Tuple defining the range of masking ratios.
+        mask_probability: Probability of applying masking.
+        dtype: Data type to cast the collated tensors.
+        n_tokens: Number of tokens for masking.
+        mask_generator: Function to generate masks.
+
+    Returns:
+        A dictionary with collated global/local crops and corresponding masks.
+    """
+    if MONAI_IMPORT_ERROR is not None:
+        raise ImportError(
+            "MONAI is required to use extended_collate_dino but not installed. "
+            "Please install MONAI to use this collate function."
+        ) from MONAI_IMPORT_ERROR
+    
+    # Apply MONAI's list_data_collate
+    collated_data = list_data_collate(samples_list)
+
+    # Extract crops
+    global_views = torch.cat(collated_data["image_global_views"], dim=0)
+    local_views = torch.cat(collated_data["image_local_views"], dim=0)
+
+    return {
+        "global_views": global_views,
+        "local_views": local_views,
+    }
+    
+
+def extended_collate_siglip(
+    samples_list: List,
+    tokenizer: Optional[Callable] = None,
+    tokenizer_padding: bool = True,
+    tokenizer_truncation: bool = True,
+    tokenizer_max_length: Optional[int] = 1024,
+    return_filenames: bool = False
+) -> dict:
+    """
+    Applies SigLIP collate and then extends it with tokenization logic.
+    
+    Args:
+        samples_list: List of samples containing 'image' and 'report'.
+        tokenizer: Tokenizer function to apply on the reports.
+    
+    Returns:
+        A dictionary with collated images and tokenized text.
+    """
+    if MONAI_IMPORT_ERROR is not None:
+        raise ImportError(
+            "MONAI is required to use extended_collate_siglip but not installed. "
+            "Please install MONAI to use this collate function."
+        ) from MONAI_IMPORT_ERROR
+    
+    collated_data = list_data_collate(samples_list)
+
+    if return_filenames:
+        if "image" in collated_data.keys():
+            if (
+                hasattr(samples_list[0]["image"].data, "meta") 
+                and "filename_or_obj" in samples_list[0]["image"].data.meta
+            ):
+                collated_data["filename"] = [s["image"].data.meta["filename_or_obj"] for s in samples_list]
+
+    if tokenizer is not None and "report" in collated_data.keys():
+        tokenizer_output = tokenizer.batch_encode_plus(
+            collated_data["report"], 
+            add_special_tokens=True,
+            padding=tokenizer_padding,
+            truncation=tokenizer_truncation,
+            max_length=tokenizer_max_length,
+        )
+        
+        collated_data["input_ids"] = torch.tensor(tokenizer_output["input_ids"])
+        collated_data["attention_mask"] = torch.tensor(tokenizer_output["attention_mask"])
+
+    return collated_data
+
+
+def collate_add_filenames(samples_list: List) -> dict:
+    """
+    Applies MONAI's list_data_collate and adds filenames to the collated output.
+
+    Args:
+        samples_list: List of samples containing 'image' with metadata.
+    Returns:
+        A dictionary with collated images and filenames.
+    """
+    if MONAI_IMPORT_ERROR is not None:
+        raise ImportError(
+            "MONAI is required to use collate_add_filenames but not installed. "
+            "Please install MONAI to use this collate function."
+        ) from MONAI_IMPORT_ERROR
+    
+    collated_data = list_data_collate(samples_list)
+
+    if "image" in collated_data.keys():
+        if (
+            hasattr(samples_list[0]["image"].data, "meta") 
+            and "filename_or_obj" in samples_list[0]["image"].data.meta
+        ):
+            collated_data["filename"] = [s["image"].data.meta["filename_or_obj"] for s in samples_list]
+
+    return collated_data

spectre/utils/config.py

@@ -0,0 +1,91 @@
+import os
+import math
+
+from spectre.utils import _utils, distributed
+
+OMEGACONF_IMPORT_ERROR = None
+try:
+    from omegaconf import OmegaConf
+except ImportError as e:
+    OmegaConf = None  # type: ignore
+    OMEGACONF_IMPORT_ERROR = e
+
+
+def apply_scaling_rules_to_cfg(cfg):
+    """
+    Apply learing rate scaling rules to the configuration object.
+    """
+    base_lr = cfg.optim.base_lr
+    cfg.optim.lr = base_lr
+
+    # Apply scaling rules
+    if cfg.optim.scaling_rule == "constant":
+        return cfg
+    
+    try:
+        scaling_type, ref_batch_size = cfg.optim.scaling_rule.split("_wrt_")
+        ref_batch_size = float(ref_batch_size)
+    except ValueError:
+        raise NotImplementedError(f"Unknown scaling rule: {cfg.optim.scaling_rule}")
+    
+    scale_factor = cfg.train.batch_size_per_gpu * distributed.get_global_size()
+    scale_factor /= ref_batch_size
+    scale_factor *= cfg.train.grad_accum_steps
+    
+    if scaling_type == "sqrt":
+        cfg.optim.lr *= math.sqrt(scale_factor)
+    elif scaling_type == "linear":
+        cfg.optim.lr *= scale_factor
+    else:
+        raise NotImplementedError(f"Unsupported scaling type: {scaling_type}")
+    
+    return cfg
+
+
+def write_config(cfg, output_dir, name="config.yaml"):
+    if OMEGACONF_IMPORT_ERROR is not None:
+        raise ImportError(
+            "OmegaConf is required to use write_config but not installed. "
+            "Please install OmegaConf to use this function."
+        ) from OMEGACONF_IMPORT_ERROR
+    
+    saved_cfg_path = os.path.join(output_dir, name)
+    with open(saved_cfg_path, "w") as f:
+        OmegaConf.save(config=cfg, f=f)
+    return saved_cfg_path
+
+
+def get_cfg_from_args(args, default_config):
+    if OMEGACONF_IMPORT_ERROR is not None:
+        raise ImportError(
+            "OmegaConf is required to use get_cfg_from_args but not installed. "
+            "Please install OmegaConf to use this function."
+        ) from OMEGACONF_IMPORT_ERROR
+    
+    args.output_dir = os.path.abspath(args.output_dir)
+    args.opts = [] if args.opts is None else args.opts
+    args.opts += [f"train.output_dir={args.output_dir}"]
+    default_cfg = OmegaConf.create(default_config)
+    cfg = OmegaConf.load(args.config_file)
+    cfg = OmegaConf.merge(default_cfg, cfg, OmegaConf.from_cli(args.opts))
+    return cfg
+
+
+def random_seed(args):
+    seed = getattr(args, "seed", 0)
+    rank = distributed.get_global_rank()
+
+    _utils.fix_random_seeds(seed + rank)
+
+
+def setup(args, default_config):
+    """
+    Create configs and perform basic setups.
+    """
+    cfg = get_cfg_from_args(args, default_config)
+    os.makedirs(args.output_dir, exist_ok=True)
+    random_seed(args)
+    accelerator = distributed.init_distributed(cfg)
+    apply_scaling_rules_to_cfg(cfg)
+    write_config(cfg, args.output_dir)
+    return cfg, accelerator

spectre/utils/dataloader.py

@@ -0,0 +1,126 @@
+from __future__ import annotations
+import os
+from typing import Union, Callable, Optional, List
+
+import torch
+from torch.utils.data import ConcatDataset
+
+MONAI_IMPORT_ERROR = None
+try:
+    import monai.data as data
+except ImportError as e:
+    data = None  # type: ignore
+    MONAI_IMPORT_ERROR = e
+
+
+
+def get_dataloader(
+    datasets: Union[str, List[str]],
+    data_dir: str,
+    include_reports: bool = False,
+    include_labels: bool = False,
+    cache_dataset: bool = False,
+    cache_dir: Optional[str] = None,
+    use_gds: bool = False,
+    transform: Optional[Callable] = None,
+    fraction: float = 1.0,
+    batch_size: int = 64,
+    num_workers: int = 4,
+    pin_memory: bool = True,
+    shuffle: bool = True,
+    collate_fn: Optional[Callable] = None,
+    drop_last: bool = True,
+    persistent_workers: bool = True,
+    use_thread: bool = False,
+) -> "DataLoader":
+    """
+    Get dataloader for training.
+    """
+    if MONAI_IMPORT_ERROR is not None:
+        raise ImportError(
+            "MONAI is required to use get_dataloader but not installed. "
+            "Please install MONAI to use this function."
+        ) from MONAI_IMPORT_ERROR
+
+    if isinstance(datasets, str):
+        datasets = [datasets]
+
+    # Validate constraints
+    if include_reports:
+        assert set(datasets).issubset({"ct_rate", "merlin", "inspect"}), \
+            "When include_reports=True, only 'ct_rate', 'merlin', and 'inspect' are allowed."
+    if include_labels:
+        assert set(datasets).issubset({"abdomen_atlas", "abdomenct_1k"}), \
+            "When include_labels=True, only 'abdomen_atlas' and 'abdomenct_1k' are allowed."
+    if use_gds:
+        assert cache_dataset, "GDS requires cache_dataset=True."
+        assert torch.cuda.is_available(), "GDS requires CUDA to be available."
+
+    # Dataset configurations
+    DATASET_CONFIGS = {
+        "ct_rate": {"folder": "CT-RATE", "base_name": "CTRate",
+            "extra": {"include_reports": include_reports}},
+        "inspect": {"folder": "INSPECT", "base_name": "Inspect",
+            "extra": {"include_reports": include_reports}},
+        "merlin": {"folder": "MERLIN", "base_name": "Merlin",
+            "extra": {"include_reports": include_reports}},
+        "nlst": {"folder": "NLST", "base_name": "Nlst"},
+        "amos": {"folder": "Amos", "base_name": "Amos"},
+        "abdomen_atlas": {"folder": "AbdomenAtlas1.0Mini", "base_name": "AbdomenAtlas",
+            "extra": {"include_labels": include_labels}},
+        "panorama": {"folder": "PANORAMA", "base_name": "Panorama"},
+        "abdomenct_1k": {"folder": "AbdomenCT-1K", "base_name": "AbdomenCT1K",
+            "extra": {"include_labels": include_labels}},
+    }
+    
+    datasets_list = []
+    for ds in datasets:
+        if ds.lower() not in DATASET_CONFIGS:
+            raise NotImplementedError(f"Dataset {ds} not implemented.")
+
+        cfg = DATASET_CONFIGS[ds.lower()]
+        folder = cfg["folder"]
+        extra_args = cfg.get("extra", {})
+
+        kwargs = {
+            "data_dir": os.path.join(data_dir, folder),
+            "transform": transform,
+            "fraction": fraction,
+            **extra_args,
+        }
+
+        base_name = cfg["base_name"]
+        class_suffix = "Dataset"
+        if cache_dataset:
+            class_suffix = "GDSDataset" if use_gds else "PersistentDataset"
+
+        class_name = f"{base_name}{class_suffix}"
+        DatasetClass = getattr(__import__("spectre.data", fromlist=[class_name]), class_name)
+
+        if cache_dataset:
+            kwargs["cache_dir"] = os.path.join(cache_dir, folder)
+            if use_gds:
+                kwargs["device"] = torch.cuda.current_device()
+
+        datasets_list.append(DatasetClass(**kwargs))
+
+    dataset = datasets_list[0] if len(datasets_list) == 1 else ConcatDataset(datasets_list)
+
+    loader_cls = getattr(data, "ThreadDataLoader" if use_thread else "DataLoader")
+    loader_kwargs = {
+        "dataset": dataset,
+        "batch_size": batch_size,
+        "num_workers": num_workers,
+        "shuffle": shuffle,
+        "drop_last": drop_last,
+    }
+
+    if not use_thread:
+        loader_kwargs.update({
+            "pin_memory": pin_memory, 
+            "persistent_workers": persistent_workers
+        })
+    if collate_fn is not None:
+        loader_kwargs["collate_fn"] = collate_fn
+
+    return loader_cls(**loader_kwargs)

spectre/utils/distributed.py

@@ -0,0 +1,92 @@
+import os
+
+import torch.distributed as dist
+
+ACCELERATE_IMPORT_ERROR = None
+try:
+    from accelerate import Accelerator, DataLoaderConfiguration
+except ImportError as e:
+    Accelerator = None  # type: ignore
+    DataLoaderConfiguration = None  # type: ignore
+    ACCELERATE_IMPORT_ERROR = e
+
+
+def is_enabled() -> bool:
+    """
+    Returns:
+        True if distributed training is enabled
+    """
+    return dist.is_available() and dist.is_initialized()
+
+
+def get_global_size() -> int:
+    """
+    Returns:
+        Number of processes in the distributed group
+    """
+    if not is_enabled():
+        return 1
+    return dist.get_world_size()
+
+
+def get_global_rank() -> int:
+    """
+    Returns:
+        The rank of the current process in the distributed group
+    """
+    if not is_enabled():
+        return 0
+    return dist.get_rank()
+
+
+def get_local_size() -> int:
+    """
+    Returns:
+        Number of processes on the current machine
+    """
+    if not is_enabled():
+        return 1
+    return int(os.environ.get("LOCAL_SIZE", 1))
+
+
+def get_local_rank() -> int:
+    """
+    Returns:
+        The rank of the current process on the current machine
+    """
+    if not is_enabled():
+        return 0
+    return int(os.environ.get("LOCAL_RANK", 0))
+
+
+def init_distributed(cfg):
+    """
+    Initialize distributed training.
+    """
+    if ACCELERATE_IMPORT_ERROR is not None:
+        raise ImportError(
+            "Accelerate is required to use init_distributed but not installed. "
+            "Please install Accelerate to use this function."
+        ) from ACCELERATE_IMPORT_ERROR
+
+    # Initialize accelerator
+    dataloader_config = DataLoaderConfiguration(
+        non_blocking=cfg.train.pin_memory,
+    )
+    accelerator = Accelerator(
+        gradient_accumulation_steps=cfg.train.grad_accum_steps,
+        log_with="wandb" if cfg.train.log_wandb else None,
+        dataloader_config=dataloader_config,
+    )
+
+    # Initialize wandb
+    if cfg.train.log_wandb:
+        accelerator.init_trackers(
+            project_name="spectre",
+            config={k: v for d in cfg.values() for k, v in d.items()},
+            init_kwargs={
+                "dir": os.path.join(cfg.train.output_dir, "logs"),
+            },
+        )
+    
+    return accelerator

spectre/utils/lora.py

@@ -0,0 +1,38 @@
+import torch.nn as nn
+import loralib as lora
+
+
+def add_lora_adapters(
+        root_module: nn.Module,
+        r: int = 8,
+        lora_alpha: int = 32,
+        lora_dropout: float = 0.05,
+        target_keywords: tuple[str, ...] = ("q_proj", "k_proj", "v_proj", "o_proj")
+    ) -> None:
+    """
+    Recursively traverses the model and replaces every `nn.Linear`
+    whose name contains one of `target_keywords` with a LoRA-augmented
+    linear layer from loralib.
+    """
+
+    for name, child in list(root_module.named_children()):
+        # If the child is itself a container, recurse first
+        add_lora_adapters(child, r, lora_alpha, lora_dropout, target_keywords)
+
+        # Replace target linear layers
+        if isinstance(child, nn.Linear) and any(k in name for k in target_keywords):
+            lora_layer = lora.Linear(                             # loralib wrapper
+                in_features=child.in_features,
+                out_features=child.out_features,
+                r=r,
+                lora_alpha=lora_alpha,
+                lora_dropout=lora_dropout,
+                bias=child.bias is not None,
+            )
+
+            # copy original weights so that behaviour is identical pre-training
+            lora_layer.weight.data = child.weight.data.clone()
+            if child.bias is not None:
+                lora_layer.bias.data = child.bias.data.clone()
+
+            setattr(root_module, name, lora_layer)               # hot-swap!

spectre/utils/masking.py

@@ -0,0 +1,196 @@
+import math
+from typing import Optional, Tuple, Union
+
+import torch
+
+
+def _random_block_mask(
+    size: Tuple[int, int, int],
+    num_masks: int,
+    min_num_masks_per_block: int = 4,
+    max_num_masks_per_block: Optional[int] = None,
+    max_attempts_per_block: int = 10,
+    generator: Optional[torch.Generator] = None,
+    device: Optional[Union[torch.device, str]] = None,
+) -> torch.Tensor:
+    """3D helper: generate a (H, W, D) boolean mask by placing cuboidal blocks.
+
+    - size: (H, W, D)
+    - num_masks: target total number of masked voxels for this image
+    - min_num_masks_per_block / max_num_masks_per_block: voxel-range per block
+    """
+    H, W, D = size
+    total = H * W * D
+    num_masks = min(max(0, int(num_masks)), total)
+
+    if max_num_masks_per_block is None:
+        max_num_masks_per_block = max(1, num_masks)
+
+    mask = torch.zeros((H, W, D), dtype=torch.bool, device=device)
+    masked_count = 0
+    global_attempts = 0
+
+    orders = [(0, 1, 2), (1, 2, 0), (2, 0, 1)]
+
+    # Try to place blocks until we have enough masked voxels or we exceed attempts
+    while masked_count < num_masks and global_attempts < max_attempts_per_block:
+        global_attempts += 1
+
+        # choose target voxels for this block
+        target_voxels = int(torch.randint(
+            min_num_masks_per_block, max_num_masks_per_block + 1, (1,), generator=generator
+        ).item())
+
+        found = False
+        local_attempts = 0
+        while not found and local_attempts < max_attempts_per_block:
+            local_attempts += 1
+
+            # random pick order for dims to reduce bias
+            order_idx = int(torch.randint(0, 3, (1,), generator=generator).item())
+            order = orders[order_idx]
+
+            # pick first dimension
+            if order[0] == 0:
+                h = int(torch.randint(1, min(H, target_voxels) + 1, (1,), generator=generator).item())
+            elif order[0] == 1:
+                w = int(torch.randint(1, min(W, target_voxels) + 1, (1,), generator=generator).item())
+            else:
+                d = int(torch.randint(1, min(D, target_voxels) + 1, (1,), generator=generator).item())
+
+            # progressively choose remaining dims while ensuring feasibility
+            try:
+                if order[0] == 0:
+                    # h chosen -> pick w then compute d_needed
+                    max_w = max(1, min(W, target_voxels // h))
+                    w = int(torch.randint(1, max_w + 1, (1,), generator=generator).item())
+                    d_needed = math.ceil(target_voxels / (h * w))
+                    if d_needed <= D:
+                        d = max(1, d_needed)
+                        found = True
+                elif order[0] == 1:
+                    # w chosen -> pick d then compute h_needed
+                    max_d = max(1, min(D, target_voxels // w))
+                    d = int(torch.randint(1, max_d + 1, (1,), generator=generator).item())
+                    h_needed = math.ceil(target_voxels / (d * w))
+                    if h_needed <= H:
+                        h = max(1, h_needed)
+                        found = True
+                else:
+                    # d chosen -> pick h then compute w_needed
+                    max_h = max(1, min(H, target_voxels // d))
+                    h = int(torch.randint(1, max_h + 1, (1,), generator=generator).item())
+                    w_needed = math.ceil(target_voxels / (d * h))
+                    if w_needed <= W:
+                        w = max(1, w_needed)
+                        found = True
+
+            except ValueError:
+                # in case of invalid ranges (defensive); just continue trying
+                continue
+
+            # fallback alternative attempt: try simple factorization heuristics
+            if not found:
+                # attempt small-to-large factorization
+                for hh in range(1, min(H, target_voxels) + 1):
+                    for ww in range(1, min(W, target_voxels // hh) + 1):
+                        dd = math.ceil(target_voxels / (hh * ww))
+                        if dd <= D:
+                            h, w, d = hh, ww, dd
+                            found = True
+                            break
+                    if found:
+                        break
+
+        if not found:
+            # couldn't find a fitting block this global attempt; move on
+            continue
+
+        # clamp block dims to volume just in case and ensure at least 1
+        h = min(max(1, int(h)), H)
+        w = min(max(1, int(w)), W)
+        d = min(max(1, int(d)), D)
+
+        # choose random location so block fits
+        x0 = int(torch.randint(0, (H - h) + 1, (1,), generator=generator).item()) if H - h > 0 else 0
+        y0 = int(torch.randint(0, (W - w) + 1, (1,), generator=generator).item()) if W - w > 0 else 0
+        z0 = int(torch.randint(0, (D - d) + 1, (1,), generator=generator).item()) if D - d > 0 else 0
+
+        mask[x0 : x0 + h, y0 : y0 + w, z0 : z0 + d] = True
+        masked_count = int(mask.sum().item())
+
+    # If still short, fill remaining voxels at random positions
+    if masked_count < num_masks:
+        remaining = num_masks - masked_count
+        indices = torch.nonzero(~mask, as_tuple=False)
+        if indices.numel() > 0:
+            perm = torch.randperm(indices.shape[0], generator=generator, device=mask.device)
+            pick = indices[perm[:remaining]]
+            mask[pick[:, 0], pick[:, 1], pick[:, 2]] = True
+
+    return mask
+
+
+def random_block_mask(
+    size: Tuple[int, int, int, int],
+    batch_mask_ratio: float = 0.5,
+    min_image_mask_ratio: float = 0.1,
+    max_image_mask_ratio: float = 0.5,
+    min_num_masks_per_block: int = 4,
+    max_num_masks_per_block: Optional[int] = None,
+    max_attempts_per_block: int = 10,
+    generator: Optional[torch.Generator] = None,
+    device: Optional[Union[torch.device, str]] = None,
+) -> torch.Tensor:
+    """Create random block masks for 3D volumes only.
+
+    Args:
+        size: (B, H, W, D)
+        batch_mask_ratio: fraction of images in the batch to apply masking to
+        min_image_mask_ratio / max_image_mask_ratio: per-image mask fraction range
+        min_num_masks_per_block / max_num_masks_per_block: voxels per block range
+        max_attempts_per_block: attempts to find a fitting block
+        generator: optional torch.Generator for reproducibility.
+        device: device for returned tensor
+
+    Returns:
+        boolean tensor with shape (B, H, W, D)
+    """
+    if len(size) != 4:
+        raise ValueError("size must be (B, H, W, D) for 3D masking.")
+
+    B, H, W, D = size
+
+    if max_image_mask_ratio < min_image_mask_ratio:
+        raise ValueError("max_image_mask_ratio must be >= min_image_mask_ratio.")
+
+    num_images_masked = int(B * batch_mask_ratio)
+    probs = torch.linspace(min_image_mask_ratio, max_image_mask_ratio, num_images_masked + 1).tolist()
+
+    image_masks = []
+    total_voxels = H * W * D
+
+    for prob_min, prob_max in zip(probs[:-1], probs[1:]):
+        # choose number of masked voxels for this image
+        u = float(prob_min + (prob_max - prob_min) * torch.rand(1, generator=generator).item())
+        num_mask = int(total_voxels * u)
+        image_masks.append(
+            _random_block_mask(
+                size=(H, W, D),
+                num_masks=num_mask,
+                min_num_masks_per_block=min_num_masks_per_block,
+                max_num_masks_per_block=max_num_masks_per_block,
+                max_attempts_per_block=max_attempts_per_block,
+                generator=generator,
+                device=device,
+            )
+        )
+
+    # Add non-masked images (all False) to fill the batch
+    for _ in range(num_images_masked, B):
+        image_masks.append(torch.zeros((H, W, D), dtype=torch.bool, device=device))
+
+    perm = torch.randperm(B, generator=generator).tolist()
+    image_masks = [image_masks[i] for i in perm]
+    
+    return torch.stack(image_masks)

spectre/utils/modeling.py

@@ -0,0 +1,550 @@
+from __future__ import annotations
+
+import math
+from enum import Enum
+from typing import List, Tuple, Optional, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+
+def deactivate_requires_grad_and_to_eval(model: nn.Module):
+    """Deactivates the requires_grad flag for all parameters of a model.
+
+    This has the same effect as permanently executing the model within a `torch.no_grad()`
+    context. Use this method to disable gradient computation and therefore
+    training for a model.
+
+    Examples:
+        >>> backbone = resnet18()
+        >>> deactivate_requires_grad(backbone)
+    """
+    for param in model.parameters():
+        param.requires_grad = False
+    model.eval()
+
+
+def activate_requires_grad_and_to_train(model: nn.Module):
+    """Activates the requires_grad flag for all parameters of a model.
+
+    Use this method to activate gradients for a model (e.g. after deactivating
+    them using `deactivate_requires_grad(...)`).
+
+    Examples:
+        >>> backbone = resnet18()
+        >>> activate_requires_grad(backbone)
+    """
+    for param in model.parameters():
+        param.requires_grad = True
+    model.train()
+
+
+@torch.no_grad()
+def update_momentum(model: nn.Module, model_ema: nn.Module, m: float):
+    """Updates parameters of `model_ema` with Exponential Moving Average of `model`
+
+    Momentum encoders are a crucial component fo models such as MoCo or BYOL.
+
+    Examples:
+        >>> backbone = resnet18()
+        >>> projection_head = MoCoProjectionHead()
+        >>> backbone_momentum = copy.deepcopy(moco)
+        >>> projection_head_momentum = copy.deepcopy(projection_head)
+        >>>
+        >>> # update momentum
+        >>> update_momentum(moco, moco_momentum, m=0.999)
+        >>> update_momentum(projection_head, projection_head_momentum, m=0.999)
+    """
+    for model_ema, model in zip(model_ema.parameters(), model.parameters()):
+        model_ema.data = model_ema.data * m + model.data * (1.0 - m)
+
+
+def update_drop_path_rate(
+    model: "VisionTransformer",
+    drop_path_rate: float,
+    mode: str = "linear",
+) -> None:
+    """Updates the drop path rate in a VisionTransformer model.
+
+    Args:
+        model:
+            VisionTransformer model.
+        drop_path_rate:
+            Maximum drop path rate.
+        mode:
+            Drop path rate update mode. Can be "linear" or "uniform". Linear increases
+            the drop path rate from 0 to drop_path_rate over the depth of the model.
+            Uniform sets the drop path rate to drop_path_rate for all blocks.
+    Raises:
+        ValueError: If an unknown mode is provided.
+    """
+    from timm.layers import DropPath
+
+    total_depth = len(model.blocks)
+
+    # Determine drop path rates based on the specified mode
+    if mode == "linear":
+        drop_probabilities = np.linspace(0, drop_path_rate, total_depth)
+    elif mode == "uniform":
+        drop_probabilities = [drop_path_rate for _ in range(total_depth)]
+    else:
+        raise ValueError(
+            f"Unknown mode: '{mode}', supported modes are 'linear' and 'uniform'."
+        )
+
+    # Update the drop path rate for each block in the model
+    for block, drop_prob in zip(model.blocks, drop_probabilities):
+        if drop_prob > 0.0:
+            block.drop_path1 = DropPath(drop_prob=drop_prob)
+            block.drop_path2 = DropPath(drop_prob=drop_prob)
+        else:
+            block.drop_path1 = nn.Identity()
+            block.drop_path2 = nn.Identity()
+
+
+def repeat_token(token: torch.Tensor, size: Tuple[int, int]) -> torch.Tensor:
+    """Repeats a token size times.
+
+    Args:
+        token: Token tensor with shape (1, 1, dim).
+        size: (batch_size, sequence_length) tuple.
+
+    Returns:
+        Tensor with shape (batch_size, sequence_length, dim) containing copies
+        of the input token.
+    """
+    batch_size, sequence_length = size
+    return token.repeat(batch_size, sequence_length, 1)
+
+
+def expand_index_like(index: torch.Tensor, tokens: torch.Tensor) -> torch.Tensor:
+    """Expands the index along the last dimension of the input tokens.
+
+    Args:
+        index:
+            Index tensor with shape (batch_size, idx_length) where each entry is
+            an index in [0, sequence_length).
+        tokens:
+            Tokens tensor with shape (batch_size, sequence_length, dim).
+
+    Returns:
+        Index tensor with shape (batch_size, idx_length, dim) where the original
+        indices are repeated dim times along the last dimension.
+    """
+    dim = tokens.shape[-1]
+    index = index.unsqueeze(-1).expand(-1, -1, dim)
+    return index
+
+
+def get_at_index(tokens: torch.Tensor, index: torch.Tensor) -> torch.Tensor:
+    """Selects tokens at index.
+
+    Args:
+        tokens:
+            Token tensor with shape (batch_size, sequence_length, dim).
+        index:
+            Index tensor with shape (batch_size, index_length) where each entry is
+            an index in [0, sequence_length).
+
+    Returns:
+        Token tensor with shape (batch_size, index_length, dim) containing the
+        selected tokens.
+    """
+    index = expand_index_like(index, tokens)
+    return torch.gather(tokens, 1, index)
+
+
+def set_at_index(
+    tokens: torch.Tensor, index: torch.Tensor, value: torch.Tensor
+) -> torch.Tensor:
+    """Copies all values into the input tensor at the given indices.
+
+    Args:
+        tokens: Tokens tensor with shape (batch_size, sequence_length, dim).
+        index: Index tensor with shape (batch_size, index_length).
+        value: Value tensor with shape (batch_size, index_length, dim).
+
+    Returns:
+        Tokens tensor with shape (batch_size, sequence_length, dim) containing
+        the new values.
+    """
+    index = expand_index_like(index, tokens)
+    return torch.scatter(tokens, 1, index, value)
+
+
+def mask_at_index(
+    tokens: torch.Tensor, index: torch.Tensor, mask_token: torch.Tensor
+) -> torch.Tensor:
+    """Copies mask token into the input tensor at the given indices.
+
+    Args:
+        tokens:
+            Tokens tensor with shape (batch_size, sequence_length, dim).
+        index:
+            Index tensor with shape (batch_size, index_length).
+        mask_token:
+            Value tensor with shape (1, 1, dim).
+
+    Returns:
+        Tokens tensor with shape (batch_size, sequence_length, dim) containing
+        the new values.
+
+    """
+    mask = tokens.new_zeros(tokens.shape)
+    mask = set_at_index(mask, index, 1)
+    return (1 - mask) * tokens + mask * mask_token
+
+
+def mask_bool(tokens: torch.Tensor, mask: torch.Tensor, mask_token: torch.Tensor) -> torch. Tensor:
+    """Returns a tensor with tokens replaced by the mask tokens in all positions where
+    the mask is True.
+
+    Args:
+        tokens:
+            Tokens tensor with shape (batch_size, sequence_length, dim).
+        mask:
+            Boolean mask tensor with shape (batch_size, sequence_length).
+        mask_token:
+            Mask token with shape (1, 1, dim).
+
+    Returns:
+        Tokens tensor with shape (batch_size, sequence_length, dim) where tokens[i, j]
+        is replaced by the mask token if mask[i, j] is True.
+    """
+    # Convert to int for multiplication.
+    mask = mask.unsqueeze(-1).to(torch.bool).to(torch.int)
+    return (1 - mask) * tokens + mask * mask_token
+
+
+def patchify(images: torch.Tensor, patch_size: Tuple[int, int, int]) -> torch.Tensor:
+    """Converts a batch of input images into patches.
+
+    Args:
+        images:
+            Images tensor with shape (batch_size, channels, height, width, depth)
+        patch_size:
+            Patch size in pixels. Image width and height must be multiples of
+            the patch size.
+
+    Returns:
+        Patches tensor with shape (batch_size, num_patches, channels * math.prod(patch_size))
+        where num_patches = image_width / patch_size * image_height / patch_size.
+
+    """
+    N, C, H, W, D = images.shape
+    assert (
+        H % patch_size[0] == 0
+        and W % patch_size[1] == 0
+        and D % patch_size[2] == 0
+    ), "Image height, width, and depth must be multiples of the patch size."
+
+    patch_h =  H // patch_size[0]
+    patch_w =  W // patch_size[1]
+    patch_d =  D // patch_size[2]
+
+    num_patches = patch_h * patch_w * patch_d
+    patches = images.reshape(shape=(
+        N, C, 
+        patch_h, patch_size[0], 
+        patch_w, patch_size[1], 
+        patch_d, patch_size[2],
+    ))
+    patches = torch.einsum("nchpwqdr->nhwdpqrc", patches)
+    patches = patches.reshape(shape=(N, num_patches, math.prod(patch_size) * C))
+    return patches
+
+
+def random_token_mask(
+    size: Tuple[int, int],
+    mask_ratio: float = 0.6,
+    mask_class_token: bool = False,
+    device: Optional[Union[torch.device, str]] = None,
+) -> torch.Tensor:
+    """Creates random token masks.
+
+    Args:
+        size:
+            Size of the token batch for which to generate masks.
+            Should be (batch_size, sequence_length).
+        mask_ratio:
+            Percentage of tokens to mask.
+        mask_class_token:
+            If False the class token is never masked. If True the class token
+            might be masked.
+        device:
+            Device on which to create the index masks.
+
+    Returns:
+        A (index_keep, index_mask) tuple where each index is a tensor.
+        index_keep contains the indices of the unmasked tokens and has shape
+        (batch_size, num_keep). index_mask contains the indices of the masked
+        tokens and has shape (batch_size, sequence_length - num_keep).
+        num_keep is equal to sequence_length * (1- mask_ratio).
+
+    """
+    batch_size, sequence_length = size
+    num_keep = int(sequence_length * (1 - mask_ratio))
+
+    noise = torch.rand(batch_size, sequence_length, device=device)
+    if not mask_class_token and sequence_length > 0:
+        # make sure that class token is not masked
+        noise[:, 0] = -1
+        num_keep = max(1, num_keep)
+
+    # get indices of tokens to keep
+    indices = torch.argsort(noise, dim=1)
+    idx_keep = indices[:, :num_keep]
+    idx_mask = indices[:, num_keep:]
+
+    return idx_keep, idx_mask
+
+
+def resample_abs_pos_embed(
+        posemb: torch.Tensor,
+        new_size: List[int],
+        old_size: List[int],
+        num_prefix_tokens: int = 1,
+        interpolation: str = 'trilinear',
+):
+    # sort out sizes, assume square if old size not provided
+    num_pos_tokens = posemb.shape[1]
+    num_new_tokens = new_size[0] * new_size[1] * new_size[2] + num_prefix_tokens
+    if num_new_tokens == num_pos_tokens and new_size[0] == new_size[1]:
+        return posemb
+
+    if num_prefix_tokens:
+        posemb_prefix, posemb = posemb[:, :num_prefix_tokens], posemb[:, num_prefix_tokens:]
+    else:
+        posemb_prefix, posemb = None, posemb
+
+    # do the interpolation
+    embed_dim = posemb.shape[-1]
+    orig_dtype = posemb.dtype
+    posemb = posemb.float()  # interpolate needs float32
+    posemb = posemb.reshape(1, old_size[0], old_size[1], old_size[2], -1).permute(0, 4, 1, 2, 3)
+    posemb = F.interpolate(posemb, size=new_size, mode=interpolation)
+    posemb = posemb.permute(0, 2, 3, 4, 1).reshape(1, -1, embed_dim)
+    posemb = posemb.to(orig_dtype)
+
+    # add back extra (class, etc) prefix tokens
+    if posemb_prefix is not None:
+        posemb = torch.cat([posemb_prefix, posemb], dim=1)
+
+    return posemb
+
+
+def resample_abs_pos_embed_nhwdc(
+        posemb: torch.Tensor,
+        new_size: List[int],
+        interpolation: str = 'trilinear',
+):
+    if new_size[0] == posemb.shape[-4] and new_size[1] == posemb.shape[-3] and new_size[2] == posemb.shape[-2]:
+        return posemb
+
+    orig_dtype = posemb.dtype
+    posemb = posemb.float()
+    posemb = posemb.reshape(1, posemb.shape[-4], posemb.shape[-3], posemb.shape[-2], posemb.shape[-1]).permute(0, 4, 1, 2, 3)
+    posemb = F.interpolate(posemb, size=new_size, mode=interpolation)
+    posemb = posemb.permute(0, 2, 3, 4, 1).to(orig_dtype)
+
+    return posemb
+
+
+def resample_patch_embed(
+        patch_embed,
+        new_size: List[int],
+        interpolation: str = 'trilinear',
+):
+    """Resample the weights of the patch embedding kernel to target resolution.
+    We resample the patch embedding kernel by approximately inverting the effect
+    of patch resizing.
+
+    Code based on:
+      https://github.com/google-research/big_vision/blob/b00544b81f8694488d5f36295aeb7972f3755ffe/big_vision/models/proj/flexi/vit.py
+
+    With this resizing, we can for example load a B/8 filter into a B/16 model
+    and, on 2x larger input image, the result will match.
+
+    Args:
+        patch_embed: original parameter to be resized.
+        new_size (tuple[int, int, int]): target shape (depth, height, width).
+        interpolation (str): interpolation for resize
+    Returns:
+        Resized patch embedding kernel.
+    """
+    import numpy as np
+    try:
+        from torch import vmap
+    except ImportError:
+        from functorch import vmap
+
+    assert len(patch_embed.shape) == 5, "Five dimensions expected"
+    assert len(new_size) == 3, "New shape should only be (height, width, depth)"
+    old_size = patch_embed.shape[-3:]
+    if tuple(old_size) == tuple(new_size):
+        return patch_embed
+
+    def resize(x_np, _new_size):
+        x_tf = torch.Tensor(x_np)[None, None, ...]
+        x_upsampled = F.interpolate(
+            x_tf, size=_new_size, mode=interpolation)[0, 0, ...].numpy()
+        return x_upsampled
+
+    def get_resize_mat(_old_size, _new_size):
+        mat = []
+        for i in range(np.prod(_old_size)):
+            basis_vec = np.zeros(_old_size)
+            basis_vec[np.unravel_index(i, _old_size)] = 1.
+            mat.append(resize(basis_vec, _new_size).reshape(-1))
+        return np.stack(mat).T
+
+    resize_mat = get_resize_mat(old_size, new_size)
+    resize_mat_pinv = torch.tensor(np.linalg.pinv(resize_mat.T), device=patch_embed.device)
+
+    def resample_kernel(kernel):
+        resampled_kernel = resize_mat_pinv @ kernel.reshape(-1)
+        return resampled_kernel.reshape(new_size)
+
+    v_resample_kernel = vmap(vmap(resample_kernel, 0, 0), 1, 1)
+    orig_dtype = patch_embed.dtype
+    patch_embed = patch_embed.float()
+    patch_embed = v_resample_kernel(patch_embed)
+    patch_embed = patch_embed.to(orig_dtype)
+    return patch_embed
+
+
+def feature_take_indices(
+        num_features: int,
+        indices: Optional[Union[int, List[int]]] = None,
+        as_set: bool = False,
+) -> Tuple[List[int], int]:
+    """ Determine the absolute feature indices to 'take' from.
+
+    Note: This function can be called in forward() so must be torchscript compatible,
+    which requires some incomplete typing and workaround hacks.
+
+    Args:
+        num_features: total number of features to select from
+        indices: indices to select,
+          None -> select all
+          int -> select last n
+          list/tuple of int -> return specified (-ve indices specify from end)
+        as_set: return as a set
+
+    Returns:
+        List (or set) of absolute (from beginning) indices, Maximum index
+    """
+    if indices is None:
+        indices = num_features  # all features if None
+
+    if isinstance(indices, int):
+        # convert int -> last n indices
+        assert 0 < indices <= num_features, f'last-n ({indices}) is out of range (1 to {num_features})'
+        take_indices = [num_features - indices + i for i in range(indices)]
+    else:
+        take_indices: List[int] = []
+        for i in indices:
+            idx = num_features + i if i < 0 else i
+            assert 0 <= idx < num_features, f'feature index {idx} is out of range (0 to {num_features - 1})'
+            take_indices.append(idx)
+
+    if not torch.jit.is_scripting() and as_set:
+        return set(take_indices), max(take_indices)
+
+    return take_indices, max(take_indices)
+
+
+def global_pool_nlc(
+        x: torch.Tensor,
+        pool_type: str = 'token',
+        num_prefix_tokens: int = 1,
+        reduce_include_prefix: bool = False,
+):
+    if not pool_type:
+        return x
+
+    if pool_type == 'token':
+        x = x[:, 0]  # class token
+    else:
+        x = x if reduce_include_prefix else x[:, num_prefix_tokens:]
+        if pool_type == 'avg':
+            x = x.mean(dim=1)
+        elif pool_type == 'avgmax':
+            x = 0.5 * (x.amax(dim=1) + x.mean(dim=1))
+        elif pool_type == 'max':
+            x = x.amax(dim=1)
+        else:
+            assert not pool_type, f'Unknown pool type {pool_type}'
+
+    return x
+
+
+def cat_keep_shapes(
+    x_list: List[torch.Tensor]
+) -> Tuple[torch.Tensor, List[Tuple[int, ...]], List[int]]:
+    if not x_list:
+        return torch.empty(0), [], []
+
+    shapes = [x.shape for x in x_list]
+    num_tokens = [x.select(dim=-1, index=0).numel() for x in x_list]
+    x_cat = torch.cat([x.flatten(0, -2) for x in x_list], dim=0)
+
+    return x_cat, shapes, num_tokens
+
+
+def uncat_with_shapes(
+    x_cat: torch.Tensor,
+    shapes: List[Tuple[int, ...]],
+    num_tokens: List[int]
+) -> List[torch.Tensor]:
+    if not shapes:
+        return []
+
+    x_splitted = torch.split_with_sizes(x_cat, num_tokens, dim=0)
+    shapes_adjusted = [shape[:-1] + torch.Size([x_cat.shape[-1]]) for shape in shapes]
+    outputs_reshape = [x.reshape(shape) for x, shape in zip(x_splitted, shapes_adjusted)]
+
+    return outputs_reshape
+
+
+def last_token_pool(
+    last_hidden_states: torch.Tensor, 
+    attention_mask: torch.Tensor
+) -> torch.Tensor:
+    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
+    if left_padding:
+        return last_hidden_states[:, -1]
+    else:
+        sequence_lengths = attention_mask.sum(dim=1) - 1
+        batch_size = last_hidden_states.shape[0]
+        return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device),
+                                  sequence_lengths]
+
+
+class Format(str, Enum):
+    NCHWD = 'NCHWD'
+    NHWDC = 'NHWDC'
+    NCL = 'NCL'
+    NLC = 'NLC'
+
+
+def nchwd_to(x: torch.Tensor, fmt: Format):
+    if fmt == Format.NHWDC:
+        x = x.permute(0, 2, 3, 4, 1)
+    elif fmt == Format.NLC:
+        x = x.flatten(2).transpose(1, 2)
+    elif fmt == Format.NCL:
+        x = x.flatten(2)
+    return x
+
+
+def nhwdc_to(x: torch.Tensor, fmt: Format):
+    if fmt == Format.NCHWD:
+        x = x.permute(0, 4, 1, 2, 3)
+    elif fmt == Format.NLC:
+        x = x.flatten(1, 2)
+    elif fmt == Format.NCL:
+        x = x.flatten(1, 2).transpose(1, 2)
+    return x

spectre/utils/param_groups.py

@@ -0,0 +1,118 @@
+import torch.nn as nn
+
+
+def get_vit_lr_decay_rate(
+    name: str,
+    llrd_factor: float = 1.0,
+    num_layers: int = 12,
+    force_is_backbone: bool = False,
+    shift: int = 0,
+) -> float:
+    """
+    Get the layer-wise learning rate decay (LLRD) rate for a given parameter name.
+
+    Args:
+        name:
+            The name of the parameter.
+        llrd_factor:
+            The decay factor for each layer.
+        num_layers:
+            The total number of layers in the model.
+        force_is_backbone:
+            If True, forces the function to treat the parameter as part of the backbone.
+        shift:
+            An integer to shift the layer ids, useful when combining multiple modules.
+
+    Returns:
+        The learning rate multiplier for the parameter.
+    """
+    layer_id = num_layers + 1
+    if name.startswith("backbone") or force_is_backbone:
+        if (
+            ".pos_embed" in name
+            or ".patch_embed" in name
+            or ".patch_proj" in name
+            or ".mask_token" in name
+            or ".cls_token" in name
+            or ".reg_token" in name
+        ):
+            layer_id = 0
+        elif ".blocks." in name:
+            layer_id = int(name[name.find(".blocks.") :].split(".")[2]) + 1 + shift
+
+    return llrd_factor ** (num_layers + 1 - layer_id)
+
+
+def get_param_groups_with_decay(
+    model: nn.Module,
+    llrd_factor: float = 1.0,
+    patch_embed_lr_mult: float = 1.0,
+    projection_head_wd_mult: float = 1.0,
+    lora_lr_factor: float = 1.0,
+    num_layers: int | None = None,
+):
+    
+    force_is_backbone = False
+    shift = 0
+    if num_layers is not None:
+        num_layers = num_layers
+    elif hasattr(model, "n_blocks"):
+        num_layers = model.n_blocks
+        force_is_backbone = True
+    elif hasattr(model, "blocks"):
+        num_layers = len(model.blocks)
+        force_is_backbone = True
+    elif hasattr(model, "backbone") and hasattr(model.backbone, "blocks"):
+        num_layers = len(model.backbone.blocks)
+    elif hasattr(model, "backbone_student") and hasattr(model.backbone_student, "blocks"):  # DINO specific
+        num_layers = len(model.backbone_student.blocks)
+    elif hasattr(model, "backbone_student") and hasattr(model.backbone_student, "vit") and hasattr(model.backbone_student.vit, "blocks"):  # DINOv2 specific
+        num_layers = len(model.backbone_student.vit.blocks)
+    elif hasattr(model, "backbone_image") and hasattr(model.backbone_image, "blocks"):  # SigLIP specific
+        if not hasattr(model, "feature_comb_image") or model.feature_comb_image is None:
+            num_layers = len(model.backbone_image.blocks)
+        else:
+            num_layers = len(model.backbone_image.blocks) + len(model.feature_comb_image.blocks)
+            shift = len(model.backbone_image.blocks)
+            force_is_backbone = True
+    else:
+        num_layers = 0
+    
+    all_param_groups = []
+    for n, p in model.named_parameters():
+        if not p.requires_grad:
+            continue
+        if not "lora_" in n:
+            s = shift if "feature_comb" in n else 0
+            llrd_rate = get_vit_lr_decay_rate(
+                n, llrd_factor, num_layers, force_is_backbone, s,
+            )
+
+            d = {
+                "name": n,
+                "params": p,
+                "lr_mult": llrd_rate,
+                "wd_mult": 1.0,
+            }
+
+            if "head" in n or "projection" in n:
+                d["wd_mult"] = projection_head_wd_mult
+
+            # No weight-decay on biases, norm parameters, layer scale gamma, learned tokens and embeddings
+            if n.endswith("bias") or "norm" in n or "gamma" in n or "fourrier_w" in n:
+                d["wd_mult"] = 0.0
+
+            if "patch_embed" in n:
+                d["lr_mult"] *= patch_embed_lr_mult
+        
+        else:
+            # LoRA parameters
+            d = {
+                "name": n,
+                "params": p,
+                "lr_mult": lora_lr_factor,
+                "wd_mult": 1.0,
+            }
+        
+        all_param_groups.append(d)
+    return all_param_groups

spectre/utils/scheduler.py

@@ -0,0 +1,236 @@
+import warnings
+from typing import Optional
+
+import numpy as np
+import torch
+
+
+def linear_warmup_schedule(
+    step: int,
+    warmup_steps: int,
+    start_value: float,
+    end_value: float,
+) -> float:
+    if warmup_steps < 0:
+        raise ValueError(f"Warmup steps {warmup_steps} can't be negative.")
+    if step < 0:
+        raise ValueError(f"Current step number {step} can't be negative.")
+    if start_value < 0:
+        raise ValueError(f"Start value {start_value} can't be negative.")
+    if end_value <= 0:
+        raise ValueError(f"End value {end_value} can't be non-positive.")
+    if start_value > end_value:
+        raise ValueError(
+            f"Start value {start_value} must be less than or equal to end value {end_value}."
+        )
+    if step < warmup_steps:
+        return start_value + step / warmup_steps * (end_value - start_value)
+    else:
+        return end_value
+    
+
+def cosine_schedule(
+    step: int,
+    max_steps: int,
+    start_value: float,
+    end_value: float,
+    period: Optional[int] = None,
+) -> float:
+    """Use cosine decay to gradually modify start_value to reach target end_value.
+
+    Args:
+        step:
+            Current step number.
+        max_steps:
+            Total number of steps.
+        start_value:
+            Starting value.
+        end_value:
+            Target value.
+        period:
+            The number of steps over which the cosine function completes a full cycle.
+            Defaults to max_steps.
+
+    Returns:
+        Cosine decay value.
+
+    """
+    if step < 0:
+        raise ValueError(f"Current step number {step} can't be negative.")
+    if max_steps < 1:
+        raise ValueError(f"Total step number {max_steps} must be >= 1.")
+    if period is None and step > max_steps:
+        warnings.warn(
+            f"Current step number {step} exceeds max_steps {max_steps}.",
+            category=RuntimeWarning,
+        )
+    if period is not None and period <= 0:
+        raise ValueError(f"Period {period} must be >= 1")
+
+    decay: float
+    if period is not None:  # "cycle" based on period, if provided
+        decay = (
+            end_value
+            - (end_value - start_value) * (np.cos(2 * np.pi * step / period) + 1) / 2
+        )
+    elif max_steps == 1:
+        # Avoid division by zero
+        decay = end_value
+    elif step == max_steps:
+        # Special case for Pytorch Lightning which updates LR scheduler also for epoch
+        # after last training epoch.
+        decay = end_value
+    else:
+        decay = (
+            end_value
+            - (end_value - start_value)
+            * (np.cos(np.pi * step / (max_steps - 1)) + 1)
+            / 2
+        )
+    return decay
+
+
+def cosine_warmup_schedule(
+    step: int,
+    max_steps: int,
+    start_value: float,
+    end_value: float,
+    warmup_steps: int,
+    warmup_start_value: float,
+    warmup_end_value: Optional[float] = None,
+    period: Optional[int] = None,
+) -> float:
+    """Use cosine decay to gradually modify start_value to reach target end_value.
+
+    Uses linear warmup for the first warmup_steps steps.
+
+    Args:
+        step:
+            Current step number.
+        max_steps:
+            Total number of steps.
+        start_value:
+            Starting value.
+        end_value:
+            Target value.
+        warmup_steps:
+            Number of steps for warmup.
+        warmup_start_value:
+            Starting value for warmup.
+        warmup_end_value:
+            Target value for warmup. Defaults to start_value.
+        period:
+            The number of steps over which the cosine function completes a full cycle.
+            Defaults to max_steps - warmup_steps.
+    Returns:
+        Cosine decay value.
+    """
+    if warmup_steps < 0:
+        raise ValueError(f"Warmup steps {warmup_steps} can't be negative.")
+    if warmup_steps > max_steps:
+        raise ValueError(f"Warmup steps {warmup_steps} must be <= max_steps.")
+    if step > max_steps:
+        warnings.warn(
+            f"Current step number {step} exceeds max_steps {max_steps}.",
+            category=RuntimeWarning,
+        )
+
+    if warmup_end_value is None:
+        warmup_end_value = start_value
+
+    if step < warmup_steps:
+        # Use step + 1 to reach warmup_end_value at end of warmup. This means that the
+        # initial warmup_start_value is skipped which is oftentimes desired when setting
+        # it to 0 as this would result in no parameter updates.
+        return (
+            warmup_start_value
+            + (warmup_end_value - warmup_start_value) * (step + 1) / warmup_steps
+        )
+    else:
+        max_steps = max_steps - warmup_steps if period is None else 1
+        return cosine_schedule(
+            step=step - warmup_steps,
+            max_steps=max_steps,
+            start_value=start_value,
+            end_value=end_value,
+            period=period,
+        )
+
+
+class CosineWarmupScheduler(torch.optim.lr_scheduler.LambdaLR):
+    """Cosine warmup scheduler for learning rate.
+
+    Args:
+        optimizer:
+            Optimizer object to schedule the learning rate.
+        warmup_epochs:
+            Number of warmup epochs or steps.
+        max_epochs:
+            Total number of training epochs or steps.
+        last_epoch:
+            The index of last epoch or step.
+        start_value:
+            Starting learning rate.
+        end_value:
+            Target learning rate.
+        verbose:
+            If True, prints a message to stdout for each update.
+        warmup_start_value:
+            Starting learning rate for warmup.
+        warmup_end_value:
+            Target learning rate for warmup. Defaults to start_value.
+
+    Note: The `epoch` arguments do not necessarily have to be epochs. Any step or index
+    can be used. The naming follows the PyTorch convention to use `epoch` for the steps
+    in the scheduler.
+    """
+
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        warmup_epochs: int,
+        max_epochs: int,
+        last_epoch: int = -1,
+        start_value: float = 1.0,
+        end_value: float = 0.001,
+        period: Optional[int] = None,
+        verbose: bool = False,
+        warmup_start_value: float = 0.0,
+        warmup_end_value: Optional[float] = None,
+    ) -> None:
+        self.warmup_epochs = warmup_epochs
+        self.max_epochs = max_epochs
+        self.start_value = start_value
+        self.end_value = end_value
+        self.period = period
+        self.warmup_start_value = warmup_start_value
+        self.warmup_end_value = warmup_end_value
+
+        super().__init__(
+            optimizer=optimizer,
+            lr_lambda=self.scale_lr,
+            last_epoch=last_epoch,
+            verbose=verbose,
+        )
+
+    def scale_lr(self, epoch: int) -> float:
+        """Scale learning rate according to the current epoch number.
+
+        Args:
+            epoch:
+                Current epoch number.
+
+        Returns:
+            Scaled learning rate.
+
+        """
+        return cosine_warmup_schedule(
+            step=epoch,
+            max_steps=self.max_epochs,
+            start_value=self.start_value,
+            end_value=self.end_value,
+            warmup_steps=self.warmup_epochs,
+            warmup_start_value=self.warmup_start_value,
+            warmup_end_value=self.warmup_end_value,
+            period=self.period,
+        )