Add pinned KRONOS2 native multiplex embeddings

CHANGELOG.md

@@ -3,6 +3,42 @@
 Notable changes to raw2features, newest first. This project follows
 [Semantic Versioning](https://semver.org).
 
+## [0.2.1] - Unreleased
+
+This focused model update adds KRONOS2 to the existing native multiplex path. The store
+layout, geometry-only `grid_hash`, and brightfield execution path are unchanged.
+
+### Models and multiplex inputs
+
+- Add the gated `kronos2` native multiplex patch encoder: a pinned, SHA-verified
+  marker-aware ViT-B/16 that produces a 768-dimensional CLS vector from ordered marker
+  stacks. The released metadata contains 288 usable marker entries, 268 of them
+  marked as pretraining markers. This path runs in fp32 and does not download its
+  BioLinkBERT text encoder. Its recorded forward contract uses fixed eight-sample
+  microbatches, including deterministic padding of a short final microbatch. The
+  optional extra installs the authors' pinned x86-64 Linux reference stack. Only the
+  required pinned custom-code modules enter an app-owned runtime snapshot, and the
+  upstream-added import path is removed after model construction.
+- Use DRAQ5 as the nuclear-segmentation fallback when a multiplex panel has no
+  recognized DAPI, Hoechst, or DNA stain. Existing panels containing DAPI retain their
+  DAPI binding, and repeated DRAQ5 acquisitions are averaged with their physical indices
+  recorded in grid identity.
+- Install the pinned xFormers reference kernel only where its release provides wheels
+  (x86-64 Linux, CPython 3.11/3.12); Python 3.13 and other platforms use KRONOS2's
+  enforced upstream PyTorch attention fallback instead of an unvalidated xFormers build.
+  Record the resolved attention backend in each model-output fingerprint so outputs
+  remain self-describing across hosts; completed arrays from either approved path remain
+  valid, while exact header/array agreement prevents partial writes being accepted.
+- Extend repeated `--marker` selection and ordering to native multiplex encoders. The
+  resolved physical indices and order are included in output fingerprints and panel
+  provenance without changing patch-grid identity.
+- Add optional `--kronos2-additional-markers` registration for novel markers. The
+  canonical consumed CSV fields, intensity statistics, and pinned BioLinkBERT dependency
+  are bound into the model contract; BioLinkBERT is fetched only when this option is used.
+- Keep native multiplex and ordinary RGB encoders in separate requests so each receives
+  the correct source-channel view. Additive runs can still write their outputs into the
+  same store when their grid geometry matches.
+
 ## [0.2.0] - 2026-07-23
 
 This release extends named-channel multiplex support with a `channelwise` strategy, adds
@@ -130,3 +166,4 @@ What 0.1.0 provides:
 
 [0.1.0]: https://github.com/CraigMyles/raw2features/releases/tag/v0.1.0
 [0.2.0]: https://github.com/CraigMyles/raw2features/compare/v0.1.0...v0.2.0
+[0.2.1]: https://github.com/CraigMyles/raw2features/compare/v0.2.0...HEAD

CITATION.cff

@@ -20,5 +20,4 @@ keywords:
   - foundation models
   - feature extraction
 license: MIT
-version: "0.2.0"
-date-released: 2026-07-23
+version: "0.2.1"

README.md

@@ -12,8 +12,9 @@ independently swappable.
 from cloud storage, and each embedding carries the metadata needed to interpret and reuse it.
 
 By analogy to `bioformats2raw` and `raw2ometiff`, but for features: point it at a raw OME-Zarr WSI,
-choose from 30+ feature extractors (UNI/UNI2, Virchow/Virchow2, CONCH, KEEP,
-GigaPath, H-optimus, Phikon, CTransPath, …; full list in [MODELS.md](docs/MODELS.md)),
+choose from 30+ feature extractors (UNI/UNI2, Virchow/Virchow2, CONCH, KRONOS/KRONOS2,
+KEEP, GigaPath, H-optimus, Phikon, CTransPath, …; full list in
+[MODELS.md](docs/MODELS.md)),
 and get back a compact, self-describing `*.embeddings.zarr` with per-patch
 coordinates such that every embedding is relocatable to the slide.
 
@@ -82,6 +83,25 @@ The same pattern covers the other gated encoders - mostly slide encoders (e.g. `
 or Drive-hosted weights). Each model's exact install is in its [`MODELS.md`](docs/MODELS.md) row
 and the matching extra's comment in `pyproject.toml`.
 
+**KRONOS2.** The gated KRONOS2 repository is loaded from a pinned Hugging Face snapshot,
+so it needs no separate git-package install:
+
+```bash
+pip install "raw2features[all,kronos2]"
+hf auth login
+```
+
+Request access to [`MahmoodLab/KRONOS2`](https://huggingface.co/MahmoodLab/KRONOS2)
+with the institutional account required by its gate. On x86-64 Linux with CPython 3.11
+or 3.12, the extra installs the authors' pinned Torch/xFormers/Transformers/timm
+reference stack; raw2features validates this path on Python 3.12. Python 3.13 and other
+platforms use the upstream PyTorch attention fallback, which is slower and may use more
+memory. The KRONOS2 extra is not available on Python 3.14 because the pinned Torch
+release does not support it. Normal use of the released marker vocabulary does not
+download BioLinkBERT; that conditional download occurs only when a complete novel-marker
+CSV is supplied. See
+[MODALITIES.md](docs/MODALITIES.md) for marker selection and registration.
+
 **Development** (from a clone, with [uv](https://docs.astral.sh/uv/)):
 
 ```bash

docs/MODALITIES.md

@@ -15,16 +15,15 @@ entry.
 | stage | brightfield | `channelwise` RGB strategy | native multiplex encoder |
 |---|---|---|---|
 | reader | `read_region` → RGB `[H,W,3]` uint8 | `read_region_channels` → native `[H,W,C]` | `read_region_channels` → native `[H,W,C]` |
-| channel identity | none | selected positional names | positional names resolved by the model |
+| channel identity | none | selected positional names | selected positional names resolved by the model |
 | segmentation | `otsu`, `canny`, … | one nuclear channel, a recognized same-stain group, or `--no-seg` | one nuclear channel, a recognized same-stain group, or `--no-seg` |
 | embedding | one RGB input per patch | one RGB input per selected marker, then mean/concat | model-specific marker stack |
 
-Panel binding happens before receipt or store completion checks. Native multiplex
-fingerprints bind the complete effective positional panel; `channelwise` fingerprints
-bind the selected physical channel identities and order (all channels when default
-selection is used). The full effective panel remains in source/panel provenance. When
-nuclear masking is enabled, the resolved physical nuclear-channel index or same-stain
-index group is also part of grid identity.
+Panel binding happens before receipt or store completion checks. Native multiplex and
+`channelwise` fingerprints bind the selected physical channel identities and order; the
+complete effective panel remains in source/panel provenance. When nuclear masking is
+enabled, the resolved physical nuclear-channel index or same-stain index group is also
+part of grid identity.
 
 ## Channel metadata
 
@@ -73,19 +72,55 @@ full CLI contract are documented in [usage.md](usage.md#rgb-encoders-on-named-ch
 ## Native multiplex encoders
 
 Native multiplex models use the same positional-panel plumbing but consume the marker
-stack directly. The current built-in example, `kronos`, maps effective source names to
-its pinned marker vocabulary, records the kept/dropped physical mapping, and binds the
-complete effective panel into its output fingerprint before resume. Installation and
-model-specific license/access details are in [MODELS.md](MODELS.md) and
-[MODEL_LICENSES.md](MODEL_LICENSES.md).
+stack directly. `kronos` and `kronos2` map effective source names to their pinned marker
+metadata, record the physical mapping, and bind the selected indices and order into the
+output fingerprint before resume. Repeated `--marker` options select and order a subset;
+without them, the source channels are offered to the model in physical order.
+Installation and model-specific license/access details are in
+[MODELS.md](MODELS.md) and [MODEL_LICENSES.md](MODEL_LICENSES.md).
 
 ```bash
 raw2features embed SLIDE.ome.zarr OUT -m kronos --mpp 0.5
+raw2features embed SLIDE.ome.zarr OUT -m kronos2 \
+  --marker CD3 --marker CD8 --marker DAPI
 ```
 
 Use `--channel-names-file` when the source lacks a complete panel, or `--no-seg` when the
 entire image should be tiled without a nuclear mask.
 
+Nuclear masking prefers recognized DAPI, Hoechst, or DNA channels. If none is present,
+it uses DRAQ5 as a fallback. Repeated acquisitions of the selected stain are averaged;
+DAPI takes precedence over DRAQ5 rather than mixing the two dyes. The resolved physical
+indices and combination are recorded in grid identity.
+
+KRONOS2 performs exact separator-insensitive marker matching and does not reuse
+KRONOSv1's biological aliases. Its released metadata contains 288 usable entries, of
+which 268 are marked as pretraining markers. Their text vectors are already stored in
+the checkpoint, so BioLinkBERT is neither downloaded nor run.
+
+The released KRONOS2 vocabulary includes DAPI and DRAQ5, but not Hoechst or
+DNA1/DNA2 labels found in many CODEX panels. raw2features does not silently map those
+stains to DAPI. A Hoechst channel may be excluded from the encoder with an explicit
+`--marker` list while still being used by nuclear segmentation, which always sees the
+complete source panel. To include it in the KRONOS2 marker stack, register it as a novel
+marker with statistics computed for the relevant dataset.
+
+For a marker outside the released metadata, pass a complete CSV through
+`--kronos2-additional-markers`. Each row requires `marker_name`, `marker_full_name`,
+`compartment`, `family`, `family_desc`, `mean`, and `std`. The statistics describe the
+marker after KRONOS2's published dtype scaling (uint8 / 255, wider unsigned integers /
+65535, or floating values / 400) and are supplied by the user; signed integers are
+rejected and raw2features does not infer cohort statistics. Novel names must not collide
+with the pinned vocabulary, while `compartment` and `family` must reuse its categories
+exactly. This optional path downloads the separately pinned BioLinkBERT model to create
+the marker text vector. The CSV content, statistics, and text-model identity become part
+of the output fingerprint.
+
+Following the published inference example, raw2features passes DRAQ5 as
+`preferred_dapi` when DAPI is absent. The upstream preprocessor therefore applies its
+DAPI normalization statistics to that selected DRAQ5 channel. This policy is recorded
+in provenance and is not configurable in v0.2.1.
+
 ## Extending multiplex support
 
 The `raw2features.multiplex_strategies` entry point separates panel/config preparation

docs/MODELS.md

@@ -59,6 +59,7 @@ read and decide on. raw2features makes no commercial-use determination.
 | `seal_conch` ⚠ | seal | 512 | 224 | 0.5 | conv. | CC-BY-NC-ND-4.0 | yes | [HF](https://huggingface.co/MahmoodLab/SEAL) · [GH](https://github.com/mahmoodlab/SEAL) · [paper](https://arxiv.org/abs/2602.14177) |
 | `seal_univ2` ⚠ | seal | 1536 | 224 | 0.5 | conv. | CC-BY-NC-ND-4.0 | yes | [HF](https://huggingface.co/MahmoodLab/SEAL) · [GH](https://github.com/mahmoodlab/SEAL) · [paper](https://arxiv.org/abs/2602.14177) |
 | `kronos` | kronos | 384 | 224 | -¹² | n/a | CC-BY-NC-ND-4.0 | yes | [HF](https://huggingface.co/MahmoodLab/KRONOS) · [GH](https://github.com/mahmoodlab/KRONOS) · [paper](https://arxiv.org/abs/2506.03373) |
+| `kronos2` | kronos2 | 768 | 256 | -¹² | n/a | CC-BY-NC-ND-4.0 | yes | [HF](https://huggingface.co/MahmoodLab/KRONOS2) |
 
 ⚠ **SEAL is experimental in v0.2.0 and is outside the exact-weight pinning
 guarantee.** raw2features pins and SHA-256 verifies the SEAL LoRA adapter, freezes
@@ -191,7 +192,8 @@ pip install flash-attn  # slide encoders only; use a wheel/build matching torch
     weights are "available for non-commercial academic purposes." Read both before use.
 11. `ctranspath` - pretrained at 20×, but the authors recommend **1.0 µm/px (10×)**
     downstream; raw2features defaults to 0.5. Override with `--mpp 1.0` to follow them.
-12. `kronos` - a multiplex (spatial-proteomics) encoder, not H&E; it has no single µm/px.
+12. `kronos` / `kronos2` - native multiplex encoders, not H&E encoders; neither has one
+    registered µm/px default.
     See `MODALITIES.md`.
 13. `kaiko_vitl` - kaiko.ai's earlier ViT-L (distinct from `midnight`); norm is symmetric
     `[0.5]` from the card, **not** the pretrained_cfg ImageNet default (the usual base-arch
@@ -258,6 +260,31 @@ back to **1.0 µm/px** unless run alongside a model that supplies one or given a
   SHA-256-verified safetensors file and does not execute the repository's remote code or
   construct its unused BERT text tower.
 - `kronos` - needs the `[kronos]` extra; multiplex only (see `MODALITIES.md`).
+- `kronos2` - needs the `[kronos2]` extra and accepted institutional access to the gated
+  [`MahmoodLab/KRONOS2`](https://huggingface.co/MahmoodLab/KRONOS2) repository. It loads
+  an allowlisted pinned custom-code snapshot into an app-owned runtime directory, verifies
+  the weights and marker table, and removes the upstream-added snapshot import path after
+  construction. It runs in fp32 and returns the 768-dimensional CLS token. The default
+  patch is 256 px;
+  an explicit alternative must be a multiple of the ViT patch size, 16. Marker matching
+  is exact after separator-insensitive normalization and does not apply KRONOSv1 aliases.
+  Upstream recommends batches of at least eight for reproducibility. raw2features fixes
+  its effective forward microbatch at eight and pads a shorter final microbatch before
+  slicing the output back to its original length. This runtime contract is recorded;
+  bitwise equality across different hardware is not promised. On x86-64 Linux with
+  CPython 3.11 or 3.12, the `[kronos2]` extra installs the authors' pinned Torch,
+  xFormers, Transformers, and timm reference stack; this path is validated with Python
+  3.12. Python 3.13 and other platforms use the upstream PyTorch attention fallback,
+  which is slower and may use more memory. The pinned Torch release does not support
+  Python 3.14. `--compile` is not supported for KRONOS2 in v0.2.1.
+  The ordinary 288-marker path uses text embeddings already stored in the checkpoint
+  and never downloads BioLinkBERT.
+  The vocabulary includes DAPI and DRAQ5 but not Hoechst or DNA1/DNA2 labels; exclude
+  those channels explicitly or register them with dataset-specific statistics. Following
+  the upstream example, a selected DRAQ5 channel uses DAPI normalization statistics when
+  DAPI is absent; this recorded policy is not configurable in v0.2.1.
+  Supplying `--kronos2-additional-markers` enables the separately pinned BioLinkBERT path
+  for complete user-defined marker metadata and statistics.
 - `ctranspath` - custom `ConvStem` patch-embed (`embedders/convstem.py`), loaded via the
   1aurent modern-timm mirror (no pinned timm fork needed).
 - `gpfm`, `retccl`, `sp22m`, `hipt`, `lunit_bt/mocov2/swav` - built from a base timm arch +

docs/MODEL_LICENSES.md

@@ -3,9 +3,12 @@
 `raw2features` (this code) is **MIT**. It downloads model weights at runtime from each
 model's source - Hugging Face, or Google Drive (by a pinned file id + verified SHA-256) for
 `chief` / `tangle` - under *your* acceptance of that model's licence/gate, and
-**redistributes no weights**. The inference recipe for each model is re-implemented in our
-own code with the source cited (`transform_source_url` in
-`src/raw2features/embedders/registry.yaml`); we vendor no model code.
+**redistributes no weights**. Most inference recipes are implemented in our own code.
+Loaders that use a model repository's custom code, including KRONOS2, execute it only
+from the recorded pinned snapshot. The source is cited through `transform_source_url` in
+`src/raw2features/embedders/registry.yaml`. KRONOS2's executable code is pinned by its
+immutable repository commit; its weights and marker table are additionally SHA-256
+verified before that code runs.
 
 **You are responsible for complying with each model's own licence where applicable.** The
 per-model `license:` field in `registry.yaml` is the machine-readable source of truth; the
@@ -51,6 +54,12 @@ table below lists each model's weights licence.
 | `plip` | MIT (project `setup.py`; no LICENSE file) | no |
 | `keep` | MIT | no |
 | `kronos` | CC-BY-NC-ND-4.0 (KRONOS, MahmoodLab; multiplex, non-commercial) | yes |
+| `kronos2` | CC-BY-NC-ND-4.0 (KRONOS2, MahmoodLab; gated non-commercial academic terms) | yes (institutional approval) |
+
+KRONOS2's optional novel-marker path also downloads
+`michiyasunaga/BioLinkBERT-large`, whose repository declares Apache-2.0. Its source,
+immutable revision, artifact hashes, and licence are recorded in the conditional
+registration contract. It is not loaded for the released KRONOS2 marker vocabulary.
 
 ## Slide encoders
 

docs/SPEC.md

@@ -17,7 +17,7 @@ The keywords MUST, SHOULD and MAY are used as in RFC 2119.
 ## Stability
 
 raw2features is alpha (`0.x`), but the package release and store schema are versioned
-independently. Package v0.2.0 continues to write `schema_version: "0.1"` because this
+independently. Package v0.2.x continues to write `schema_version: "0.1"` because this
 release does not change the established read contract or layout. To let readers depend on
 the format now, the stable parts are scoped explicitly: