KatherLab · mducducd · Mar 31, 2026 · Mar 23, 2026 · Mar 23, 2026 · Mar 23, 2026
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
@@ -12,7 +12,7 @@ jobs:
     # needs: format_and_lint
     strategy:
       matrix:
-        python-version: ["3.11", "3.12"]
+        python-version: ["3.13"]
         extractor: [
           "ctranspath",
           "chief-ctranspath",
@@ -59,7 +59,7 @@ jobs:
     # needs: format_and_lint
     strategy:
       matrix:
-        python-version: ["3.11", "3.12"]
+        python-version: ["3.13"]
 
     steps:
     - uses: actions/checkout@v4
@@ -95,12 +95,12 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v5
       with:
-        python-version: '3.11'
+        python-version: '3.13'
     - name: Install the code linting and formatting tool Ruff
       run: pipx install ruff
     - name: Lint code with Ruff
-      run: ruff check --output-format=github --target-version=py311
+      run: ruff check --output-format=github --target-version=py313
     - name: Check code formatting with Ruff
-      run: ruff format --diff --target-version=py311
+      run: ruff format --diff --target-version=py313
       # continue-on-error: true
 
diff --git a/.python-version b/.python-version
@@ -1 +1 @@
-3.12
+3.13
diff --git a/README.md b/README.md
@@ -11,8 +11,6 @@ STAMP is an **end‑to‑end, weakly‑supervised deep‑learning pipeline** tha
 
 **Want to start now?** [Jump to Installation](#installation) or [walk through our Getting Started guide](getting-started.md) for a hands-on tutorial.
 
-
-
 ## **Why choose STAMP?**
 
 * 🚀 **Scalable**: Run locally or on HPC (SLURM) with the same CLI; built to handle multi‑center cohorts and large WSI collections.  
@@ -64,20 +62,42 @@ cd STAMP
 ```
 
 ```bash
-# GPU (CUDA) Installation (Using flash-attn on CUDA systems for gigapath and other models)
+# GPU (CUDA) Installation (excluding conchv1_5, gigapath and musk)
 
-# And then this for all models:
-uv sync --extra build --extra gpu
+uv sync --extra gpu
 source .venv/bin/activate
 ```
 
 ```bash
-# CPU-only Installation (excluding COBRA, Gigapath (and flash-attn))
+# CPU-only Installation (excluding conchv1_5, gigapath and musk)
 
 uv sync --extra cpu
 source .venv/bin/activate
 ```
 
+> [!CAUTION]
+> In the next step we will build [flash-attn](https://github.com/dao-ailab/flash-attention), this might take an extended amount of time and consume a lot of RAM and CPU time!
+>
+> Please make sure you have [Nvidia CUDA Toolkit 13.0](https://developer.nvidia.com/cuda-13-0-2-download-archive) installed! You must use Nvidia Driver version 580 or newer!
+>
+> The `nvcc --version` command must indicate that 13.0 is installed and is currently in PATH: `Cuda compilation tools, release 13.0, V13.0.88`.
+>
+> If you get another version or `Command 'nvcc' not found`, add it to the PATH:
+> ```bash
+> export CUDA_HOME=/usr/local/cuda-13.0
+> export PATH="${CUDA_HOME}/bin:$PATH"
+> ```
+>
+> Run `nvcc --version` to ensure flash-attn will be built for CUDA 13.0
+
+
+```bash
+# GPU (CUDA) Installation - building flash-attn for supporting conchv1_5, gigapath and musk
+
+MAX_JOBS=2 uv sync --extra gpu_all # to speed up the build time increase max_jobs! This might use more RAM!
+source .venv/bin/activate
+```
+
 If you encounter errors during installation please read Installation Troubleshooting [below](#installation-troubleshooting).
 
 ### Additional Dependencies
@@ -204,7 +224,6 @@ uv sync --extra build --extra gpu
 ```
 
 ## Reproducibility
-
-We use a central `Seed` utility to set seeds for PyTorch, NumPy, and Python’s `random`. This makes data loading and model initialization reproducible. Always call `Seed.set(seed)` once at startup.
-
-We do not enable [`torch.use_deterministic_algorithms()`](https://pytorch.org/docs/stable/notes/randomness.html#reproducibility) because it can cause large performance drops. Expect runs with the same seed to follow the same training trajectory, but not bit-for-bit identical low-level kernels.
+> [!NOTE]
+> We use a central `Seed` utility to set seeds for PyTorch, NumPy, and Python’s `random`. This makes data loading and model initialization reproducible. Always call `Seed.set(seed)` once at startup.
+> We do not enable [`torch.use_deterministic_algorithms()`](https://pytorch.org/docs/stable/notes/randomness.html#reproducibility) because it can cause large performance drops. Expect runs with the same seed to follow the same training trajectory, but not bit-for-bit identical low-level kernels.
diff --git a/getting-started.md b/getting-started.md
@@ -7,6 +7,9 @@ you will need some WSIs,
 a table mapping each of these slides to a patient
 as well as some ground truth we will eventually train a neural network on.
 
+> [!NOTE]
+> If you prefer a browser-based workflow, see [`STAMP-Workbench`](https://github.com/KatherLab/STAMP-Workbench), a separate web UI for building and monitoring STAMP pipelines. You can install it into the same environment with `uv pip install git+https://github.com/KatherLab/STAMP-Workbench` and launch it from your STAMP checkout with `stamp-workbench`.
+
 ## Whole Slide Images
 
 The whole slide images have to be in any of the formats [supported by OpenSlide][openslide].
@@ -410,7 +413,6 @@ The key differences for patient-level modeling are:
 You can then run statistics as done with tile-level features.
 
 ## Heatmaps and Top Tiles
-<img src="docs/overlay-heatmap.png" width="500px" align="center"></img>
 
 The `stamp heatmaps` command generates visualization outputs to help interpret model predictions and identify which regions of the slide contribute most to the classification decision. This command creates:
 
@@ -495,7 +497,7 @@ This section lets you control global training parameters, model type, and the ta
 
 advanced_config:
   seed: 42
-  task: "classification" # or regression/survial
+  task: "classification" # or regression/survival
   max_epochs: 32
   patience: 16
   batch_size: 64
@@ -522,9 +524,11 @@ advanced_config:
       use_alibi: false
 ```
 
-STAMP automatically adapts its **model architecture**, **loss function**, and **evaluation metrics** based on the task specified in the configuration file.
-
-**Regression** tasks only require `ground_truth_label`.  
-**Survival analysis** tasks require `time_label` (follow-up time) and `status_label` (event indicator).  
-**Multi-target classification** requires `ground_truth_label` as a list and `advanced_config.model_name: "barspoon"`.
-These requirements apply consistently across cross-validation, training, deployment, and statistics.
+> [!NOTE]
+> STAMP automatically adapts its **model architecture**, **loss function**, and **evaluation metrics** based on the task specified in the configuration file.
+>
+> - **Regression**: requires only `ground_truth_label`.
+> - **Survival analysis**: requires `time_label` (follow-up time) and `status_label` (event indicator).
+> - **Multi-target classification**: requires `ground_truth_label` as a list and `advanced_config.model_name: "barspoon"`.
+>
+> These requirements apply consistently across cross-validation, training, deployment, and statistics.
diff --git a/pyproject.toml b/pyproject.toml
@@ -14,7 +14,7 @@ authors = [
 ]
 description = "A protocol for Solid Tumor Associative Modeling in Pathology"
 readme = "README.md"
-requires-python = ">=3.11"
+requires-python = ">=3.13,<3.14"
 
 classifiers = [
     "Programming Language :: Python :: 3",
@@ -23,31 +23,30 @@ classifiers = [
 ]
 
 dependencies = [
-    "beartype>=0.21.0",
-    "einops>=0.8.1",
-    "h5py>=3.14.0",
-    "jaxtyping>=0.3.2",
-    "lightning>=2.5.2",
-    "matplotlib>=3.10.5",
-    "numpy>=2.2.2",
-    "opencv-python>=4.12.0.88",
-    "openpyxl>=3.1.5",
-    "openslide-bin>=4.0.0.8",
-    "openslide-python>=1.4.2",
-    "packaging>=25.0",
-    "pandas>=2.3.1",
-    "pillow>=11.3.0",
-    "pydantic>=2.11.7",
-    "pyyaml>=6.0.2",
-    "scikit-learn>=1.6.1",
-    "scipy>=1.16.1",
-    "torch>=2.7.1,<2.8.0",
-    "torchmetrics>=1.6.0",
-    "torchvision>=0.22.1",
-    "tqdm>=4.67.1",
-    "timm>=1.0.19",
-    "transformers>=4.55.0",
-    "lifelines>=0.28.0",
+    "beartype~=0.22.9",
+    "einops~=0.8.2",
+    "h5py~=3.16.0",
+    "jaxtyping~=0.3.9",
+    "lightning~=2.6.1",
+    "matplotlib~=3.10.8",
+    "numpy~=2.4.3",
+    "opencv-python~=4.13.0.92",
+    "openpyxl~=3.1.5",
+    "openslide-bin~=4.0.0.13",
+    "openslide-python~=1.4.3",
+    "packaging~=26.0",
+    "pandas~=2.3.3",
+    "pillow~=12.1.1",
+    "pydantic~=2.12.5",
+    "pyyaml~=6.0.3",
+    "scikit-learn~=1.8.0",
+    "scipy~=1.17.1",
+    "torchmetrics~=1.9.0",
+    "tqdm~=4.67.3",
+    "timm~=1.0.25",
+    "transformers~=4.57.6",
+    "lifelines~=0.30.3",
+    "huggingface-hub~=0.36.2"
 ]
 
 [project.optional-dependencies]
@@ -58,14 +57,13 @@ build = [
     "ninja"
 ]
 flash-attention = [
-    "flash-attn>=2.8.3",
+    "flash-attn==2.8.3",
     ]
 conch = [
     "huggingface-hub>=0.26.2",
     "conch @ git+https://github.com/KatherLab/CONCH",
 ]
 conch1_5_cpu = [
-    "transformers>=4.45.2",
     "einops-exts==0.0.4",
 ]
 conch1_5 = [
@@ -99,19 +97,21 @@ virchow2 = [
     "torch>=2.0.0",
 ]
 cobra = [
-    "stamp[flash-attention]",
-    "causal-conv1d>=1.5.3.post1",
-    "mamba-ssm>=2.2.6.post3",
+    "causal-conv1d @ https://github.com/Dao-AILab/causal-conv1d/releases/download/v1.6.1.post4/causal_conv1d-1.6.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl ; sys_platform == 'linux' and platform_machine == 'x86_64' and python_version == '3.13' and extra == 'cobra'",
+    "causal-conv1d @ https://github.com/Dao-AILab/causal-conv1d/releases/download/v1.6.1.post4/causal_conv1d-1.6.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_aarch64.whl ; sys_platform == 'linux' and platform_machine == 'aarch64' and python_version == '3.13' and extra == 'cobra'",
+    "mamba-ssm @ https://github.com/state-spaces/mamba/releases/download/v2.3.1/mamba_ssm-2.3.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl ; sys_platform == 'linux' and platform_machine == 'x86_64' and python_version == '3.13'",
+    "mamba-ssm @ https://github.com/state-spaces/mamba/releases/download/v2.3.1/mamba_ssm-2.3.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_aarch64.whl ; sys_platform == 'linux' and platform_machine == 'aarch64' and python_version == '3.13'",
+
     "cobra @ git+http://github.com/KatherLab/COBRA.git@49d231191db5a9c2ea37a2398dd922c8c9ee9cdb",
-    "jinja2>=3.1.4", 
-    "triton"
+    "jinja2>=3.1.4",
+    "triton",
 ]
 prism_cpu = [
     "sacremoses==0.1.1",
     "environs==11.0.0",
 ]
 prism = [
-    "stamp[prism_cpu, flash-attention]",
+    "stamp[prism_cpu]",
 ]
 madeleine = [
     "madeleine @ git+https://github.com/mahmoodlab/MADELEINE.git@de7c85acc2bdad352e6df8eee5694f8b6f288012"
@@ -127,16 +127,29 @@ plip_cpu = [
     "transformers>=4.45.2"
 ]
 plip = [
-    "stamp[plip_cpu, flash-attention]",
+    "stamp[plip_cpu]",
 ]
 mcp = [
     "fastmcp>=2.11.2",
 ]
 
 
 # Blanket target
-cpu = ["stamp[conch,ctranspath,uni,virchow2,chief_ctranspath,conch1_5_cpu,prism_cpu,madeleine,musk_cpu,plip_cpu]"]
-gpu = ["stamp[conch,ctranspath,uni,virchow2,chief_ctranspath,conch1_5,prism,madeleine,musk,plip,gigapath,cobra]"]
+cpu = [
+    "torch~=2.10.0",
+    "torchvision~=0.25.0",
+    "stamp[conch,ctranspath,uni,virchow2,chief_ctranspath,conch1_5_cpu,prism_cpu,madeleine,musk_cpu,plip_cpu]"
+]
+gpu = [
+    "torch~=2.10.0",
+    "torchvision~=0.25.0",
+    "stamp[conch,ctranspath,uni,virchow2,chief_ctranspath,conch1_5_cpu,prism,madeleine,plip,musk_cpu,cobra]"
+]
+gpu_all = [
+    "torch~=2.10.0",
+    "torchvision~=0.25.0",
+    "stamp[conch,ctranspath,uni,virchow2,chief_ctranspath,musk,gigapath,conch1_5,prism,madeleine,plip,cobra]"
+]
 all = ["stamp[cpu]"]
 
 [project.scripts]
@@ -148,11 +161,11 @@ all = ["stamp[cpu]"]
 
 [dependency-groups]
 dev = [
-    "huggingface-hub>=0.34.3",
-    "ipykernel>=6.30.1",
-    "pyright>=1.1.403",
-    "pytest>=8.4.1",
-    "ruff>=0.12.7",
+    "huggingface-hub~=0.36.2",
+    "ipykernel~=7.2.0",
+    "pyright~=1.1.408",
+    "pytest~=9.0.2",
+    "ruff>=0.15.7",
 ]
 
 [build-system]
@@ -173,28 +186,41 @@ lint.ignore = ["F722"] # https://docs.kidger.site/jaxtyping/faq/#flake8-or-ruff-
 
 [tool.uv]
 conflicts = [
-    [
-        { extra = "cpu" },
-        { extra = "gpu" }
-    ]
+  [
+    { extra = "cpu" },
+    { extra = "gpu" },
+  ],
+  [
+    { extra = "cpu" },
+    { extra = "gpu_all" },
+  ],
 ]
 build-constraint-dependencies = [
-  "torch<2.8",
-  "torchvision<0.23",
+  "torch==2.10.0",
+  "torchvision==0.25.0",
 ]
 
 
 [tool.uv.sources]
 torch = [
-  { index = "pytorch-cu128", marker = "sys_platform != 'darwin'" }
+  { index = "pytorch-cpu", extra = "cpu" },
+  { index = "pytorch-cu130", extra = "gpu" },
+  { index = "pytorch-cu130", extra = "gpu_all" },
 ]
 torchvision = [
-  { index = "pytorch-cu128", marker = "sys_platform != 'darwin'" }
+  { index = "pytorch-cpu", extra = "cpu" },
+  { index = "pytorch-cu130", extra = "gpu" },
+  { index = "pytorch-cu130", extra = "gpu_all" },
 ]
 
 [[tool.uv.index]]
-name = "pytorch-cu128"
-url = "https://download.pytorch.org/whl/cu128"
+name = "pytorch-cpu"
+url = "https://download.pytorch.org/whl/cpu"
+explicit = true
+
+[[tool.uv.index]]
+name = "pytorch-cu130"
+url = "https://download.pytorch.org/whl/cu130"
 explicit = true
 
 
@@ -210,7 +236,7 @@ requires-dist = [
     "scikit-learn",
     "tqdm",
     "transformers",
-    "xformers; sys_platform != 'darwin'"  # xformers is not supported on macOS
+    "xformers; (sys_platform == 'linux' and platform_machine == 'x86_64') or (sys_platform == 'win32' and platform_machine == 'AMD64')"
 ]
 
 
@@ -221,23 +247,7 @@ requires-dist = [
     "einops",
 ]
 
-
-[[tool.uv.dependency-metadata]]
-name = "mamba-ssm"
-requires-dist = [
-    "setuptools",
-]
-
-[[tool.uv.dependency-metadata]]
-name = "causal-conv1d"
-requires-dist = [
-    "setuptools",
-    "torch"
-]
-
 [tool.uv.extra-build-dependencies]
-cobra = [{ requirement = "torch", match-runtime = true }]
 flash-attn = [{ requirement = "torch", match-runtime = true }]
-mamba-ssm = [{ requirement = "torch", match-runtime = true }]
 gigapath = [{ requirement = "torch", match-runtime = true }]
 conch = [{ requirement = "torch", match-runtime = true }]