From ebddfd30f474adfaa57db757a1a60d27e1695d5d Mon Sep 17 00:00:00 2001 From: Craig Myles Date: Wed, 29 Jul 2026 13:10:36 +0100 Subject: [PATCH 1/4] Add pinned KRONOS2 native multiplex embeddings --- CHANGELOG.md | 27 + CITATION.cff | 4 +- README.md | 22 +- docs/MODALITIES.md | 54 +- docs/MODELS.md | 27 +- docs/MODEL_LICENSES.md | 15 +- docs/SPEC.md | 2 +- docs/usage.md | 91 +- pyproject.toml | 26 +- src/raw2features/__init__.py | 2 +- src/raw2features/cli/_validation.py | 18 + src/raw2features/cli/embed.py | 25 +- src/raw2features/cli/embed_many.py | 25 +- src/raw2features/cli/verify.py | 22 +- src/raw2features/embedders/base.py | 128 ++- src/raw2features/embedders/fingerprint.py | 54 ++ .../embedders/kronos2_embedder.py | 882 +++++++++++++++++ .../embedders/kronos2_metadata.py | 362 +++++++ src/raw2features/embedders/registry.yaml | 36 + src/raw2features/multiplex/__init__.py | 3 + src/raw2features/multiplex/panel.py | 73 ++ src/raw2features/pipeline/runner.py | 448 ++++++++- tests/test_channelwise_runner.py | 2 +- tests/test_cli_commands.py | 51 + tests/test_kronos2.py | 910 ++++++++++++++++++ tests/test_kronos2_metadata.py | 274 ++++++ tests/test_native_multiplex_selection.py | 366 +++++++ tests/test_slide_embedders.py | 2 +- uv.lock | 386 +++----- 29 files changed, 4017 insertions(+), 320 deletions(-) create mode 100644 src/raw2features/embedders/kronos2_embedder.py create mode 100644 src/raw2features/embedders/kronos2_metadata.py create mode 100644 src/raw2features/multiplex/panel.py create mode 100644 tests/test_kronos2.py create mode 100644 tests/test_kronos2_metadata.py create mode 100644 tests/test_native_multiplex_selection.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 32750ee..745d221 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,32 @@ 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. +- 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 +156,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 diff --git a/CITATION.cff b/CITATION.cff index c294b44..12c2fa4 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -20,5 +20,5 @@ keywords: - foundation models - feature extraction license: MIT -version: "0.2.0" -date-released: 2026-07-23 +version: "0.2.1" +date-released: 2026-07-29 diff --git a/README.md b/README.md index b53053f..5d799c6 100644 --- a/README.md +++ b/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,23 @@ 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, the extra installs +the authors' pinned Torch/xFormers/Transformers/timm reference stack; raw2features +validates this path on Python 3.12. 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 diff --git a/docs/MODALITIES.md b/docs/MODALITIES.md index 3e13f1c..7e34a84 100644 --- a/docs/MODALITIES.md +++ b/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,50 @@ 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. +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 diff --git a/docs/MODELS.md b/docs/MODELS.md index 7113904..ec2a768 100644 --- a/docs/MODELS.md +++ b/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,29 @@ 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. The `[kronos2]` extra + installs the authors' pinned Torch, xFormers, Transformers, and timm reference stack + on x86-64 Linux and is validated with Python 3.12. 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 + diff --git a/docs/MODEL_LICENSES.md b/docs/MODEL_LICENSES.md index 1ec2cb4..2afbac8 100644 --- a/docs/MODEL_LICENSES.md +++ b/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 diff --git a/docs/SPEC.md b/docs/SPEC.md index 2ade9ce..a356886 100644 --- a/docs/SPEC.md +++ b/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: diff --git a/docs/usage.md b/docs/usage.md index 84d0775..bbe87b5 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -139,19 +139,23 @@ Any token passed on the command line is redacted from stored command provenance. | `-m, --model` | `resnet50` | model name; **repeatable** for multi-model (`-f`/`--feature-extractor` are aliases) | | `--mpp` | per-model | target µm/px; default = each model's **recommended** MPP (0.5 for most pathology FMs, 1.0 for scale-agnostic). A value overrides the scale while retaining each model's recommended patch size; combine it with `--patch-size` to force one shared grid | | `--source-mpp` | source metadata | level-0 source µm/px override for an uncalibrated OME-Zarr whose spatial axes declare no physical unit; normally leave unset | -| `--patch-size` | per-model | patch side in px; default = each model's recommended size (224 / 448 CONCH / 512 conch_v1_5). A value overrides the size; combine it with `--mpp` to force one shared grid | +| `--patch-size` | per-model | patch side in px; default = each model's recommended size (224 for most / 256 KRONOS2 / 448 CONCH / 512 conch_v1_5). A value overrides the size; combine it with `--mpp` to force one shared grid | | `--step` | = patch | stride in output px; `< patch` gives overlap | | `--no-seg` | off | tile the whole slide (skip tissue masking) | | `--segmenter` | `otsu` | tissue segmenter to use | | `--tissue-threshold` | `0.1` | keep a grid cell if ≥ this fraction is tissue | | `--features-dtype` | `float16` | stored feature dtype | | `--stain-norm` | off | stain-normalize each patch before embedding (`macenko`\|`reinhard`\|`vahadane`; `vahadane` needs `raw2features[stain]`); **changes the features** - use separate output dirs for with/without-norm experiments | +| `--multiplex-strategy` | none | adapt an ordinary RGB model to a named-channel multiplex source; built-in: `channelwise` | +| `--marker` | all source channels | select a multiplex marker; repeat to set the subset and order for a native model or strategy | +| `--channel-names-file` | source metadata | complete positional `.txt`, `.csv`, or `.tsv` panel when OME channel labels are absent or incomplete | +| `--kronos2-additional-markers` | none | complete KRONOS2 novel-marker CSV; conditionally loads the pinned BioLinkBERT text encoder | | `--config FILE` | none | YAML extraction plan (`extractions:` list of `{model, mpp?, patch_px?}`); supersedes `-m`/`--mpp`/`--patch-size`. The same model may repeat (one grid each - the MPP-ablation case) | | `-s, --slide-encoder` | none | slide-level encoder(s) run after patch embedding (e.g. `titan`); reads patch features from the store, no WSI re-read. Repeatable. | | `--qc` | none | per-patch QC producer(s) writing `grids//qc//` (e.g. `grandqc`); needs the producer's extra (`raw2features[grandqc]`) | | `--qc-stain-norm` | off | normalize the QC input first (`macenko`\|`reinhard`\|`vahadane`; `vahadane` needs `raw2features[stain]`) - for a stain outside the QC model's domain | | `--device` | `auto` | `auto` (best of cuda→mps→cpu) \| `cuda` \| `mps` \| `cpu` | -| `--devices` | = `--device` | opt-in in-process multi-GPU, e.g. `cuda:0,cuda:1`; shards this slide's patches across them and gathers in coord order (output identical to one device). Multi-device execution is not supported with `--multiplex-strategy`; see _In-process multi-GPU_ in §6. | +| `--devices` | = `--device` | opt-in in-process multi-GPU, e.g. `cuda:0,cuda:1`; shards this slide's patches across them and gathers in coord order (output identical to one device). Patch-parallel multi-device execution is not supported with multiplex strategies or native multiplex models; see _In-process multi-GPU_ in §6. | | `--batch-size` | `256` | patches per forward pass (lower it on small GPUs) | | `--amp` | `auto` | precision: `auto` (each model's card precision) \| `fp32` \| `bf16` \| `fp16` | | `--snap-to-level` | off | read a pyramid level natively (no resample); MPP = that level | @@ -320,6 +324,75 @@ the same panel-binding and fingerprint contract. Third-party strategies can acce JSON object through `--multiplex-params`; the built-in `channelwise` strategy uses its explicit options and does not accept additional strategy parameters. +### Native multiplex encoders + +Native multiplex encoders consume selected marker channels together rather than +converting each channel to RGB. Run KRONOS2 after installing its optional stack and +accepting the gated institutional-access terms: + +```bash +pip install "raw2features[all,kronos2]" +hf auth login + +raw2features embed multiplex_slide.ome.zarr multiplex_out \ + -m kronos2 --marker CD3 --marker CD8 --marker DAPI +``` + +`kronos2` is a marker-aware ViT-B/16 that returns a 768-dimensional CLS vector in fp32. +Its default patch is 256 pixels. A custom `--patch-size` is accepted only when it is a +multiple of 16. Marker names are matched exactly after the model's +separator-insensitive normalization; KRONOSv1 aliases and compound-name fallbacks are +not applied. `--marker` selects physical channels and fixes their order before +completion checks. The full source panel, selected indices, model-resolved names, and +order remain in provenance and the output fingerprint. + +The released vocabulary contains DAPI and DRAQ5 but not Hoechst or DNA1/DNA2 labels. +Those common CODEX nuclear stains are not silently aliased. Exclude one from the +encoder with an explicit `--marker` list while still allowing nuclear segmentation to +use it, or register it through `--kronos2-additional-markers` with dataset-specific +statistics. When DAPI is absent and DRAQ5 is selected, raw2features follows the upstream +example by passing DRAQ5 as `preferred_dapi`, which applies the model's DAPI +normalization statistics; the recorded policy is not configurable in v0.2.1. + +Upstream recommends batches of at least eight for reproducibility. raw2features uses +fixed eight-sample forward microbatches and pads a shorter final microbatch before +discarding the padded outputs. This runtime contract is recorded. The global +`--batch-size` may be left unchanged; it does not override this internal cap. Bitwise +equality across different hardware is not promised. The `[kronos2]` extra installs the +authors' pinned x86-64 Linux reference stack and is validated with Python 3.12. The +pinned Torch release does not support Python 3.14. `--compile` is not supported for +KRONOS2 in v0.2.1; use the validated eager path. + +The released KRONOS2 repository includes 288 usable marker-metadata entries, 268 marked +as pretraining markers, and stores their text vectors in the checkpoint. Normal +inference therefore never downloads BioLinkBERT. To register a marker outside that +released metadata, supply a complete CSV: + +```bash +raw2features embed multiplex_slide.ome.zarr multiplex_out \ + -m kronos2 --marker NovelMarker --marker DAPI \ + --kronos2-additional-markers additional_markers.csv +``` + +Every CSV row requires `marker_name`, `marker_full_name`, `compartment`, `family`, +`family_desc`, `mean`, and `std`. The mean and standard deviation must be finite, with a +positive standard deviation. `compartment` and `family` must exactly reuse categories +from KRONOS2's pinned marker metadata, and a novel name must not collide with that +vocabulary under the model's separator-insensitive matching. + +Statistics are supplied in the scaled intensity domain consumed by KRONOS2: uint8 +values divided by 255, wider unsigned integers divided by 65535, or floating values +divided by 400. Signed-integer sources are rejected. raw2features does not estimate +these dataset statistics. The option loads the separately pinned BioLinkBERT model, +about 1.33 GB, only for novel-marker text vectors. The canonical CSV content, intensity +statistics, BioLinkBERT revision, and text-encoding recipe are included in the model +fingerprint. + +Run native multiplex and ordinary RGB encoders in separate commands so each receives the +correct source-channel representation. An additive second command can still append its +features to an existing matching grid. Native multiplex patch execution uses one device +per slide; `embed-many --devices` can distribute different slides across devices. + ### Segmentation The default `otsu` segmenter runs Otsu-on-saturation + morphology on a low-res @@ -345,9 +418,9 @@ devices in one process; it never changes the embeddings (it is excluded from the config hash), so a store built with it is identical to a single-device one and resumes interchangeably. -Multi-device `--devices` execution does not currently support `--multiplex-strategy`. -Strategy-derived multiplex runs use one device per slide; prefer `--device` for that -selection. +Multi-device patch-parallel `embed --devices` execution does not currently support +multiplex strategies or native multiplex models. A multiplex run uses one device per +slide. `embed-many --devices` may still distribute distinct slides across devices. - **`embed` (one slide) - patch-parallel, latency mode.** `embed … --devices cuda:0,cuda:1` shards the slide's patches into contiguous blocks, embeds each @@ -434,9 +507,11 @@ reader, models (order-independent), segmenter / --no-seg, --mpp, --patch-size, --stain-norm, --snap-to-level, --mpp-tolerance, --allow-upsample, --amp ``` -For multiplex requests, the hash also covers the strategy and its parameters, ordered -marker selection, normalization recipe and pyramid-level limit, aggregation, the resolved -positional channel names and the resolved physical nuclear-channel indices when masking +For multiplex requests, the hash also covers the ordered physical marker selection and +resolved positional channel names. Strategy-derived requests additionally cover their +normalization recipe, pyramid-level limit, aggregation, and other strategy parameters. +Native model parameters, including KRONOS2 additional-marker metadata, are covered when +present. The resolved physical nuclear-channel indices enter grid identity when masking is enabled. The `--channel-names-file` path and serialization are not hashed; their parsed effective names are. diff --git a/pyproject.toml b/pyproject.toml index 46a620e..b15e335 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "raw2features" -version = "0.2.0" +version = "0.2.1" description = "Read OME-Zarr whole-slide images and emit patch- and slide-level foundation-model embeddings - storage backend and model independently swappable." readme = "README.md" authors = [{ name = "Craig Myles", email = "me@craig.im" }] @@ -74,6 +74,7 @@ conch_v1_5 = "raw2features.embedders.conch_v1_5_embedder:ConchV1_5Embedder" transformers = "raw2features.embedders.transformers_embedder:TransformersEmbedder" dino_teacher = "raw2features.embedders.dino_teacher_embedder:DinoTeacherEmbedder" kronos = "raw2features.embedders.kronos_embedder:KronosEmbedder" +kronos2 = "raw2features.embedders.kronos2_embedder:Kronos2Embedder" musk = "raw2features.embedders.musk_embedder:MuskEmbedder" open_clip = "raw2features.embedders.open_clip_embedder:OpenClipEmbedder" clip_hf = "raw2features.embedders.clip_hf_embedder:ClipHFEmbedder" @@ -102,9 +103,9 @@ chief = "raw2features.slide_embedders.chief:CHIEFSlideEmbedder" tangle = "raw2features.slide_embedders.tangle:TangleSlideEmbedder" [project.optional-dependencies] -# Convenience bundle: the full stack (OME-Zarr reads + segmentation + torch + -# models) in one install - `pip install "raw2features[all]"`. Excludes the git-package -# encoders (conch/kronos) and the export-only extras (spatialdata/h5), which stay opt-in. +# Convenience bundle: the general OME-Zarr, image, torch, and model stack in one install +# - `pip install "raw2features[all]"`. Special/gated model stacks (including +# conch/kronos/kronos2) and export-only extras (spatialdata/h5) stay opt-in. all = ["raw2features[zarr,image,torch,models,stain]"] # OME-Zarr reading: ngff-zarr/ome-zarr for metadata, zarr-python 3 for hot reads. zarr = ["zarr>=3.1", "numcodecs>=0.15", "ngff-zarr>=0.13", "ome-zarr>=0.10", "fsspec[http]>=2024.0"] @@ -139,6 +140,23 @@ hibou = ["transformers<5"] # KRONOS multiplex (spatial proteomics): CODEX TIFF-conversion deps (the v1 git package # is installed separately - see above). kronos = ["raw2features[torch,models]", "tifffile>=2024", "imagecodecs"] +# KRONOS2 is distributed as a gated Transformers custom-code snapshot. Its authors +# publish this exact stack for reproducible inference, so the model-specific extra +# deliberately pins it rather than silently substituting newer attention kernels. +# xFormers is part of that reference path on Linux; the upstream code has a slower, +# more memory-intensive fallback elsewhere. BioLinkBERT is not loaded for published +# markers. The reference stack is validated on Python 3.12; its pinned Torch release +# does not provide Python 3.14 wheels. +kronos2 = [ + "torch==2.6.0", + "xformers==0.0.29.post3; sys_platform == 'linux' and platform_machine == 'x86_64'", + "transformers==4.56.0", + "timm==1.0.19", + "huggingface-hub>=0.28.1,<1.0", + "omegaconf>=2.3", + "pandas>=2.2.2", + "safetensors>=0.4", +] # MUSK vision-language (BEiT3) image tower: its runtime imports (the git package itself # declares no deps, so it can't downgrade torch/timm; install it separately - see above). musk = ["raw2features[torch,models]", "fairscale", "open_clip_torch"] diff --git a/src/raw2features/__init__.py b/src/raw2features/__init__.py index 7cbdf2b..12e259d 100644 --- a/src/raw2features/__init__.py +++ b/src/raw2features/__init__.py @@ -26,7 +26,7 @@ from typing import TYPE_CHECKING -__version__ = "0.2.0" +__version__ = "0.2.1" __all__ = [ "__version__", diff --git a/src/raw2features/cli/_validation.py b/src/raw2features/cli/_validation.py index f060e45..7b3930a 100644 --- a/src/raw2features/cli/_validation.py +++ b/src/raw2features/cli/_validation.py @@ -177,3 +177,21 @@ def parse_json_object(value: str | None, param_hint: str) -> dict[str, Any]: "must be a JSON object with string keys", param_hint=param_hint ) return parsed + + +def parse_kronos2_additional_markers_file(path: str | None) -> dict[str, Any] | None: + """Parse the optional KRONOS2 marker table into its semantic model contract.""" + + if path is None: + return None + from raw2features.embedders.kronos2_metadata import ( + Kronos2MarkerMetadataError, + parse_kronos2_additional_markers, + ) + + try: + return parse_kronos2_additional_markers(path) + except (OSError, Kronos2MarkerMetadataError) as exc: + raise typer.BadParameter( + str(exc), param_hint="--kronos2-additional-markers" + ) from exc diff --git a/src/raw2features/cli/embed.py b/src/raw2features/cli/embed.py index 3c7e256..7f51df2 100644 --- a/src/raw2features/cli/embed.py +++ b/src/raw2features/cli/embed.py @@ -14,6 +14,7 @@ from ._validation import ( parse_channel_names_file, parse_json_object, + parse_kronos2_additional_markers_file, validate_amp, validate_batch_size, validate_geometry, @@ -77,8 +78,9 @@ def embed( multiplex_markers: list[str] = typer.Option( [], "--marker", - help="Multiplex marker to include; repeat in the required order. Default: all " - "named channels in source order (concat requires an explicit list).", + help="Multiplex marker to include; repeat in the required order. Applies to " + "native multiplex models and --multiplex-strategy. Default: all source " + "channels (channelwise concat requires an explicit list).", ), channel_names_file: str | None = typer.Option( None, @@ -87,6 +89,12 @@ def embed( "position. Supplies missing labels and verifies existing labels; combine " "with --marker to select or order a subset.", ), + kronos2_additional_markers: str | None = typer.Option( + None, + "--kronos2-additional-markers", + help="Complete CSV metadata for explicitly selected markers outside KRONOS2's " + "released vocabulary. Conditionally loads the pinned BioLinkBERT encoder.", + ), multiplex_normalization: str = typer.Option( "percentile", "--multiplex-normalization", @@ -233,6 +241,9 @@ def embed( validate_multiplex_percentiles(multiplex_percentile_low, multiplex_percentile_high) strategy_params = parse_json_object(multiplex_params, "--multiplex-params") channel_names_override = parse_channel_names_file(channel_names_file) + kronos2_marker_contract = parse_kronos2_additional_markers_file( + kronos2_additional_markers + ) if hf_token: os.environ["HF_TOKEN"] = hf_token os.environ["HUGGING_FACE_HUB_TOKEN"] = hf_token @@ -248,6 +259,11 @@ def embed( models = list(dict.fromkeys(e["model"] for e in geometry_config)) else: models = list(feature_extractor) + if kronos2_marker_contract is not None and "kronos2" not in models: + raise typer.BadParameter( + "requires `kronos2` in the requested model set", + param_hint="--kronos2-additional-markers", + ) # Resolve and preview the per-model extraction geometry (one grid per group). groups = resolve_geometry(models, mpp, patch_size, geometry_config) @@ -288,6 +304,11 @@ def embed( multiplex_normalization_max_side_px=multiplex_normalization_max_side_px, multiplex_aggregation=multiplex_aggregation, multiplex_strategy_params=strategy_params, + native_multiplex_params=( + {"kronos2": {"additional_markers": kronos2_marker_contract}} + if kronos2_marker_contract is not None + else {} + ), channel_names_override=channel_names_override, snap_to_level=snap_to_level, mpp_tolerance=mpp_tolerance, diff --git a/src/raw2features/cli/embed_many.py b/src/raw2features/cli/embed_many.py index 1fddf4e..34d6075 100644 --- a/src/raw2features/cli/embed_many.py +++ b/src/raw2features/cli/embed_many.py @@ -41,6 +41,7 @@ from ._validation import ( parse_channel_names_file, parse_json_object, + parse_kronos2_additional_markers_file, validate_amp, validate_batch_size, validate_geometry, @@ -111,8 +112,9 @@ def embed_many( multiplex_markers: list[str] = typer.Option( [], "--marker", - help="Multiplex marker to include; repeat in the required order. Default: all " - "named channels in source order (concat requires an explicit list).", + help="Multiplex marker to include; repeat in the required order. Applies to " + "native multiplex models and --multiplex-strategy. Default: all source " + "channels (channelwise concat requires an explicit list).", ), channel_names_file: str | None = typer.Option( None, @@ -121,6 +123,12 @@ def embed_many( "position. Applies to every slide in this command and verifies any existing " "labels; combine with --marker to select/order a subset.", ), + kronos2_additional_markers: str | None = typer.Option( + None, + "--kronos2-additional-markers", + help="Complete CSV metadata for explicitly selected markers outside KRONOS2's " + "released vocabulary. Applies to every slide in this command.", + ), multiplex_normalization: str = typer.Option( "percentile", "--multiplex-normalization", @@ -249,6 +257,9 @@ def embed_many( validate_multiplex_percentiles(multiplex_percentile_low, multiplex_percentile_high) strategy_params = parse_json_object(multiplex_params, "--multiplex-params") channel_names_override = parse_channel_names_file(channel_names_file) + kronos2_marker_contract = parse_kronos2_additional_markers_file( + kronos2_additional_markers + ) from raw2features.slide_embedders.model_registry import ( validate_slide_encoder_names, ) @@ -323,6 +334,11 @@ def embed_many( models = list(dict.fromkeys(e["model"] for e in geometry_config)) else: models = list(feature_extractor) + if kronos2_marker_contract is not None and "kronos2" not in models: + raise typer.BadParameter( + "requires `kronos2` in the requested model set", + param_hint="--kronos2-additional-markers", + ) from raw2features.embedders.model_registry import build_embedder, resolve_geometry @@ -351,6 +367,11 @@ def embed_many( multiplex_normalization_max_side_px=multiplex_normalization_max_side_px, multiplex_aggregation=multiplex_aggregation, multiplex_strategy_params=strategy_params, + native_multiplex_params=( + {"kronos2": {"additional_markers": kronos2_marker_contract}} + if kronos2_marker_contract is not None + else {} + ), channel_names_override=channel_names_override, slide_encoders=list(slide_encoder), qc=list(qc), diff --git a/src/raw2features/cli/verify.py b/src/raw2features/cli/verify.py index 39538f9..f3c40b9 100644 --- a/src/raw2features/cli/verify.py +++ b/src/raw2features/cli/verify.py @@ -26,6 +26,7 @@ from ._validation import ( parse_channel_names_file, parse_json_object, + parse_kronos2_additional_markers_file, validate_amp, validate_geometry, validate_multiplex_percentiles, @@ -72,7 +73,8 @@ def verify( multiplex_markers: list[str] = typer.Option( [], "--marker", - help="Multiplex marker; repeat in the exact order used by the embed run.", + help="Multiplex marker; repeat in the exact order used by the embed run. " + "Applies to native multiplex models and --multiplex-strategy.", ), channel_names_file: str | None = typer.Option( None, @@ -80,6 +82,11 @@ def verify( help="UTF-8 .txt/.csv/.tsv with exactly one ordered name per physical C-axis " "position; must resolve to the same effective panel used by embed.", ), + kronos2_additional_markers: str | None = typer.Option( + None, + "--kronos2-additional-markers", + help="Complete KRONOS2 novel-marker CSV used by the embed run.", + ), multiplex_normalization: str = typer.Option( "percentile", "--multiplex-normalization", @@ -134,6 +141,9 @@ def verify( validate_multiplex_percentiles(multiplex_percentile_low, multiplex_percentile_high) strategy_params = parse_json_object(multiplex_params, "--multiplex-params") channel_names_override = parse_channel_names_file(channel_names_file) + kronos2_marker_contract = parse_kronos2_additional_markers_file( + kronos2_additional_markers + ) from raw2features.core.device import resolve_device from raw2features.embedders.model_registry import resolve_geometry from raw2features.pipeline.runner import resolve_run @@ -146,6 +156,11 @@ def verify( models = list(dict.fromkeys(e["model"] for e in geometry_config)) else: models = list(feature_extractor) + if kronos2_marker_contract is not None and "kronos2" not in models: + raise typer.BadParameter( + "requires `kronos2` in the requested model set", + param_hint="--kronos2-additional-markers", + ) # Verification has no programmatic external-embedder instance from which to # derive a current contract or geometry. Fail closed with the established clean # diagnostic before the geometry resolver rejects the unknown name. @@ -185,6 +200,11 @@ def verify( multiplex_normalization_max_side_px=multiplex_normalization_max_side_px, multiplex_aggregation=multiplex_aggregation, multiplex_strategy_params=strategy_params, + native_multiplex_params=( + {"kronos2": {"additional_markers": kronos2_marker_contract}} + if kronos2_marker_contract is not None + else {} + ), channel_names_override=channel_names_override, snap_to_level=snap_to_level, mpp_tolerance=mpp_tolerance, diff --git a/src/raw2features/embedders/base.py b/src/raw2features/embedders/base.py index a69829e..972d474 100644 --- a/src/raw2features/embedders/base.py +++ b/src/raw2features/embedders/base.py @@ -190,10 +190,123 @@ def set_panel(self, channel_names: list[str] | None) -> dict: """ return {} + def bind_multiplex_panel( + self, + channel_names: list[str] | None, + *, + selected_channels: list[dict[str, Any]] | None = None, + model_params: dict[str, Any] | None = None, + ) -> dict: + """Bind one native model to an ordered subset of a physical marker panel. + + ``set_panel`` is the original compatibility seam and expects its patch tensor + and names to have matching, zero-based channel positions. This wrapper keeps + that API intact: it records the requested physical indices, passes only the + selected names (in requested order) to ``set_panel``, and + :meth:`select_multiplex_channels` applies the same ordered slice to each full + source patch before the model transform runs. + + Native embedders with additional model-specific panel configuration may + override this method. The default refuses non-empty ``model_params`` so a + misspelled or unsupported option can never be silently ignored. + """ + + if model_params: + raise ValueError( + f"{self.name} does not accept native multiplex model parameters: " + f"{sorted(model_params)}" + ) + panel = ["" if value is None else str(value) for value in (channel_names or [])] + if selected_channels is None: + selected_channels = [ + {"source_index": index, "source_name": name} + for index, name in enumerate(panel) + ] + + indices: list[int] = [] + names: list[str] = [] + for record in selected_channels: + index = record.get("source_index") + name = record.get("source_name") + if ( + isinstance(index, bool) + or not isinstance(index, int) + or not 0 <= index < len(panel) + or not isinstance(name, str) + or panel[index] != name + ): + raise ValueError( + f"{self.name} received a marker selection that does not match " + "the physical source panel" + ) + indices.append(index) + names.append(name) + if len(set(indices)) != len(indices): + raise ValueError( + f"{self.name} received duplicate physical marker-channel indices" + ) + + self._raw2features_multiplex_source_indices = tuple(indices) + summary = self.set_panel(names) + if not isinstance(summary, dict): + raise TypeError(f"{self.name}.set_panel() must return a provenance mapping") + if indices != list(range(len(panel))): + # Existing native adapters report indices relative to the names passed to + # set_panel. Remap those provenance records to the physical source C axis; + # their internal tensor indices deliberately remain local to the sliced + # patch. + summary = dict(summary) + mapping = summary.get("mapping") + if isinstance(mapping, list): + remapped = [] + for row in mapping: + item = dict(row) + local_index = item.get("channel_index") + if ( + isinstance(local_index, int) + and not isinstance(local_index, bool) + and 0 <= local_index < len(indices) + ): + item["channel_index"] = indices[local_index] + remapped.append(item) + summary["mapping"] = remapped + summary["source_selection"] = [ + {"source_index": index, "source_name": name} + for index, name in zip(indices, names, strict=True) + ] + return summary + + def select_multiplex_channels( + self, patches_hwc: list[np.ndarray] + ) -> list[np.ndarray]: + """Apply the physical ordered selection bound by + :meth:`bind_multiplex_panel`. + + The reader continues to return complete HWC patches so nuclear segmentation + and other native models see the full source panel. Selection happens only at + this per-model transform boundary. + """ + + indices = getattr(self, "_raw2features_multiplex_source_indices", None) + if indices is None: + return patches_hwc + return [np.asarray(patch)[..., list(indices)] for patch in patches_hwc] + @property def embedding_dim(self) -> int: return self.spec.embedding_dim + @property + def max_batch_size(self) -> int | None: + """Optional runtime cap for one forward/transform batch. + + Most embedders leave this unset. Models whose architecture pads or expands + each sample to a fixed internal batch can override it; the runner then lowers + the user's configured batch size without changing any output identity. + """ + + return None + @property def transform_signature(self) -> tuple: """Hashable identity of this model's preprocessing. @@ -205,7 +318,7 @@ def transform_signature(self) -> tuple: ``input_size`` (which determines whether/how a resize happens), the normalisation ``mean``/``std``, and the ``interpolation`` used on resize. """ - return ( + signature = ( self.spec.input_size, self.spec.mean, self.spec.std, @@ -213,6 +326,19 @@ def transform_signature(self) -> tuple: self.spec.crop_pct, self.spec.crop_mode, ) + if self.modality == "multiplex": + # Native marker-aware encoders do not share preprocessing merely because + # the RGB-shaped registry fields happen to match. Their panel vocabulary, + # selected channels, and transform implementation are model-specific. + return ( + "native_multiplex", + self.name, + tuple( + getattr(self, "_raw2features_multiplex_source_indices", ()) or () + ), + *signature, + ) + return signature @property def transform_input_dtype(self) -> str: diff --git a/src/raw2features/embedders/fingerprint.py b/src/raw2features/embedders/fingerprint.py index 5b45374..e082a92 100644 --- a/src/raw2features/embedders/fingerprint.py +++ b/src/raw2features/embedders/fingerprint.py @@ -297,6 +297,7 @@ def _effective_patch_checkpoint(spec: ModelSpec) -> dict[str, Any]: "transformers_snapshot_return_conch", ), "kronos": (repo, "pinned_local_file"), + "kronos2": (repo, "pinned_local_snapshot"), "musk": (repo, "pinned_local_file"), "open_clip": (repo, "pinned_local_snapshot"), "keep": (repo, "pinned_safetensors_local_image_wrapper"), @@ -377,6 +378,14 @@ def _seal_composite(spec: ModelSpec) -> dict[str, Any]: def _patch_constructor(spec: ModelSpec) -> dict[str, Any]: + from .kronos2_embedder import ( + KRONOS2_FORWARD_CONTRACT, + KRONOS2_MARKER_MATCHING_CONTRACT, + KRONOS2_MARKER_METADATA_FILENAME, + KRONOS2_SNAPSHOT_ALLOW_PATTERNS, + KRONOS2_SNAPSHOT_TOP_LEVEL_ALLOWLIST, + ) + constructor: dict[str, Any] = { "timm_kwargs": deepcopy(spec.timm_kwargs), "checkpoint_load": deepcopy(spec.checkpoint), @@ -407,6 +416,34 @@ def _patch_constructor(spec: ModelSpec) -> dict[str, Any]: "marker_metadata": "marker_metadata.csv", "construction_package_revision": KRONOS_PACKAGE_REVISION, }, + "kronos2": { + "entrypoint": "transformers.AutoModel.from_pretrained", + "input": "app_owned_allowlisted_pinned_local_snapshot", + "snapshot_allow_patterns": list(KRONOS2_SNAPSHOT_ALLOW_PATTERNS), + "snapshot_top_level_allowlist": list( + KRONOS2_SNAPSHOT_TOP_LEVEL_ALLOWLIST + ), + "trust_remote_code": True, + "transformers_local_files_only": True, + "upstream_nested_download_guard": "local_directory_input", + "upstream_sys_path_cleanup": "remove_new_entry_after_construction", + "marker_metadata": { + "filename": KRONOS2_MARKER_METADATA_FILENAME, + "sha256": spec.timm_kwargs.get("marker_metadata_sha256"), + }, + "marker_matching": deepcopy(KRONOS2_MARKER_MATCHING_CONTRACT), + "preferred_dapi_policy": "canonical_dapi_else_draq5_else_none", + "published_marker_text_embeddings": "checkpoint_buffer", + "novel_marker_registration": ( + "explicit_additional_markers_parameter_with_conditional_contract" + ), + "forward": { + **deepcopy(KRONOS2_FORWARD_CONTRACT), + "method": "__call__", + "marker_names": "canonical_marker_metadata_names", + "result": "x_norm_clstoken", + }, + }, "musk": { "architecture": "musk_large_patch16_384", "with_head": False, @@ -529,12 +566,29 @@ def patch_output_fingerprint(spec: ModelSpec, resolved_amp: str) -> dict[str, An "crop_mode": spec.crop_mode, } ) + if spec.family == "kronos2": + from .kronos2_embedder import KRONOS2_SCALING_CONTRACT + + payload["preprocessing"]["native_multiplex"] = deepcopy( + KRONOS2_SCALING_CONTRACT + ) return make_output_fingerprint(payload) def resolved_patch_amp(spec: ModelSpec, requested_amp: str, device: str) -> str: """Resolve requested/card AMP to the precision the forward path really uses.""" + # KRONOS2's published inference contract is fixed fp32. Fail rather than silently + # coercing an explicit override: the raw requested AMP participates in receipt + # identity, so accepting an override that never executes would make completion + # semantics misleading. + if spec.family == "kronos2": + if requested_amp not in {"auto", "fp32"}: + raise ValueError( + "KRONOS2 follows its published fp32 inference path and does not " + f"accept requested AMP {requested_amp!r}; use 'auto' or 'fp32'" + ) + return "fp32" selected = spec.inference_amp if requested_amp == "auto" else requested_amp # Embedder._forward_ctx enables fp16/bf16 autocast only on CUDA. Passing either # dtype on CPU/MPS still executes the model in fp32, so provenance must say fp32. diff --git a/src/raw2features/embedders/kronos2_embedder.py b/src/raw2features/embedders/kronos2_embedder.py new file mode 100644 index 0000000..300b022 --- /dev/null +++ b/src/raw2features/embedders/kronos2_embedder.py @@ -0,0 +1,882 @@ +"""KRONOS2 native multiplex patch encoder. + +KRONOS2 is distributed as a gated Hugging Face repository containing custom +``transformers`` code, a marker-aware DINOv2 ViT-B/16 checkpoint, and the marker +metadata that defines both normalisation and marker identity. raw2features downloads +one immutable snapshot, verifies both content-affecting data artifacts, and gives only +that local directory to ``AutoModel``. This prevents the pinned custom loader from +silently resolving a newer copy of its own weights. + +Published markers use text embeddings already baked into the checkpoint. The upstream +``register_additional_markers`` path runs BioLinkBERT for genuinely novel markers. +raw2features enables it only when the user supplies a complete, fingerprinted marker +table. BioLinkBERT is then downloaded at an immutable revision, every model and +tokenizer artifact is verified, and the pinned upstream prompt/encoder path runs +locally in CPU float32. Ordinary published-marker inference never downloads or +instantiates that additional 1.3 GB model. +""" + +from __future__ import annotations + +import csv +import gc +import hashlib +import importlib +import json +import math +import os +import sys +import tempfile +import unicodedata +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import numpy as np + +from raw2features.core.plugins import register + +from ._hub import ( + download_pinned_hf_snapshot, + pinned_model_cache_dir, + verify_sha256, +) +from .base import Embedder +from .kronos2_metadata import ( + KRONOS2_ADDITIONAL_MARKER_COLUMNS, + KRONOS2_ADDITIONAL_MARKER_CONTRACT_VERSION, + KRONOS2_BIOLINKBERT_ARTIFACT_SHA256, + KRONOS2_BIOLINKBERT_REVISION, + KRONOS2_BIOLINKBERT_SOURCE, + KRONOS2_NOVEL_MARKER_REGISTRATION_CONTRACT, + materialize_kronos2_additional_markers_csv, +) + +if TYPE_CHECKING: # pragma: no cover + import torch + + +KRONOS2_MARKER_METADATA_FILENAME = "marker_metadata.csv" +KRONOS2_SNAPSHOT_TOP_LEVEL_ALLOWLIST = ( + "config.json", + "configuration_kronos2.py", + "dinov2", + KRONOS2_MARKER_METADATA_FILENAME, + "marker_utils.py", + "modeling_kronos2.py", + "kronos2_vitb16_teacher.pth", +) +KRONOS2_SNAPSHOT_ALLOW_PATTERNS = ( + "config.json", + "configuration_kronos2.py", + "dinov2/**", + KRONOS2_MARKER_METADATA_FILENAME, + "marker_utils.py", + "modeling_kronos2.py", + "kronos2_vitb16_teacher.pth", +) +KRONOS2_SCALING_CONTRACT: dict[str, Any] = { + "stage_1_dtype_divisor": { + "uint8": 255.0, + "other_unsigned_integer": 65535.0, + "floating": 400.0, + }, + "stage_2": "upstream_model.preprocess_per_marker_zscore", + "arithmetic_dtype": "float32", +} +KRONOS2_FORWARD_CONTRACT: dict[str, Any] = { + "backbone": { + "architecture": "dinov2_vit_base", + "patch_stride_px": 16, + "register_tokens": 16, + "positional_reference_size_px": 224, + }, + "reference_inference_patch_size_px": 256, + "accepted_patch_shape": "square_hwc_positive_side_multiple_of_16", + "upstream_patch_stride_px": 16, + "runner_max_batch_size": 8, + "transform_staging_device": "cpu", + "forward_microbatch_size": 8, + "h2d_transfer": "one_forward_microbatch_at_a_time", + "short_microbatch_padding": { + "method": "repeat_last_sample", + "target_rows": 8, + "output": "slice_back_to_unpadded_row_count", + }, + "output_order": "input_order", +} +KRONOS2_MARKER_MATCHING_CONTRACT: dict[str, Any] = { + "name": "upstream_separator_insensitive_exact", + "cleaning": { + "lowercase": True, + "translate_to_underscore": ["-", " ", ":", "(", ")"], + "drop": ["/"], + "translate": {"α": "a"}, + "strip": True, + }, + "match_key_drop": ["_", "."], + "aliases": False, + "compound_token_fallback": False, + "forward_names": "canonical_marker_metadata_names", +} + +# Re-exported here because fingerprint construction already obtains KRONOS2's +# loader contracts from this module. The conditional value itself lives beside +# the canonical additional-marker parser so CLI, runner, and loader share one +# identity rather than independently spelling the BioLinkBERT recipe. +KRONOS2_ADDITIONAL_MARKER_REGISTRATION_CONTRACT = ( + KRONOS2_NOVEL_MARKER_REGISTRATION_CONTRACT +) + + +def _canonical_json(value: Any) -> str: + return json.dumps( + value, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + allow_nan=False, + ) + + +def _validate_additional_marker_contract( + value: Any, +) -> tuple[list[dict[str, Any]], str]: + """Validate the exact semantic contract emitted by ``kronos2_metadata``.""" + + expected_keys = { + "contract_version", + "consumed_columns", + "registration", + "rows", + "content_sha256", + } + if not isinstance(value, dict) or set(value) != expected_keys: + raise ValueError( + "KRONOS2 additional_markers must be the canonical parsed contract " + f"with exactly these fields: {sorted(expected_keys)}" + ) + if value["contract_version"] != KRONOS2_ADDITIONAL_MARKER_CONTRACT_VERSION: + raise ValueError("KRONOS2 additional_markers has an unknown contract version") + if value["consumed_columns"] != list(KRONOS2_ADDITIONAL_MARKER_COLUMNS): + raise ValueError( + "KRONOS2 additional_markers has a different consumed-column contract" + ) + if value["registration"] != KRONOS2_NOVEL_MARKER_REGISTRATION_CONTRACT: + raise ValueError( + "KRONOS2 additional_markers has a different BioLinkBERT registration " + "contract" + ) + + rows = value["rows"] + if not isinstance(rows, list) or not rows: + raise ValueError("KRONOS2 additional_markers must contain at least one row") + expected_row_keys = set(KRONOS2_ADDITIONAL_MARKER_COLUMNS) + seen_casefold: set[str] = set() + seen_match_keys: set[str] = set() + canonical_rows: list[dict[str, Any]] = [] + for row_number, row in enumerate(rows, start=2): + if not isinstance(row, dict) or set(row) != expected_row_keys: + raise ValueError( + "KRONOS2 additional_markers row " + f"{row_number} must contain exactly " + f"{list(KRONOS2_ADDITIONAL_MARKER_COLUMNS)}" + ) + canonical: dict[str, Any] = {} + for column in KRONOS2_ADDITIONAL_MARKER_COLUMNS[:5]: + item = row[column] + if ( + not isinstance(item, str) + or not item + or item != item.strip() + or any(unicodedata.category(character) == "Cc" for character in item) + ): + raise ValueError( + "KRONOS2 additional_markers row " + f"{row_number} column {column!r} is not canonical text" + ) + canonical[column] = item + for column in ("mean", "std"): + item = row[column] + if not isinstance(item, float): + raise ValueError( + "KRONOS2 additional_markers row " + f"{row_number} column {column!r} must be a canonical float" + ) + number = item + if not math.isfinite(number) or (column == "std" and number <= 0): + raise ValueError( + "KRONOS2 additional_markers row " + f"{row_number} column {column!r} is invalid" + ) + canonical[column] = 0.0 if number == 0 else number + + marker_name = canonical["marker_name"] + casefold_key = marker_name.casefold() + match_key = _marker_match_key(marker_name) + if casefold_key in seen_casefold or match_key in seen_match_keys: + raise ValueError( + "KRONOS2 additional_markers contains duplicate marker identities" + ) + seen_casefold.add(casefold_key) + seen_match_keys.add(match_key) + canonical_rows.append(canonical) + + semantic = { + "contract_version": KRONOS2_ADDITIONAL_MARKER_CONTRACT_VERSION, + "consumed_columns": list(KRONOS2_ADDITIONAL_MARKER_COLUMNS), + "registration": KRONOS2_NOVEL_MARKER_REGISTRATION_CONTRACT, + "rows": canonical_rows, + } + digest = hashlib.sha256(_canonical_json(semantic).encode("utf-8")).hexdigest() + if value["content_sha256"] != digest: + raise ValueError( + "KRONOS2 additional_markers content_sha256 does not match its " + "canonical rows and registration contract" + ) + return canonical_rows, digest + + +def _clean_marker_name(name: str) -> str: + """Mirror the immutable upstream ``clean_marker_name`` contract.""" + + return ( + str(name) + .lower() + .translate( + str.maketrans( + { + "-": "_", + " ": "_", + ":": "_", + "α": "a", + "(": "_", + ")": "_", + "/": "", + } + ) + ) + .strip() + ) + + +def _marker_match_key(name: str) -> str: + """Return the exact separator-insensitive key used by pinned KRONOS2.""" + + return _clean_marker_name(name).replace("_", "").replace(".", "") + + +def _scaling_factor(dtype: np.dtype) -> float: + """Return the upstream raw-pixel divisor for one multiplex source dtype.""" + + dtype = np.dtype(dtype) + if dtype == np.dtype(np.uint8): + return 255.0 + if np.issubdtype(dtype, np.unsignedinteger): + return 65535.0 + if np.issubdtype(dtype, np.floating): + return 400.0 + raise ValueError( + f"KRONOS2 has no published input scaling for dtype {dtype!r}; " + "expected uint8, a wider unsigned integer, or floating point" + ) + + +def _metadata_sha256(spec) -> str: + value = spec.timm_kwargs.get("marker_metadata_sha256") + if not isinstance(value, str) or len(value) != 64: + raise ValueError( + f"{spec.name}: registry timm_kwargs.marker_metadata_sha256 must be " + "a 64-character SHA-256 digest" + ) + return value + + +def _assert_allowlisted_runtime_snapshot(snapshot: str) -> None: + """Refuse a cached runtime directory containing unrequested top-level code.""" + + allowed = { + *KRONOS2_SNAPSHOT_TOP_LEVEL_ALLOWLIST, + ".cache", # huggingface_hub local-directory download metadata + "__pycache__", # may appear after a previous successful local load + } + actual = {path.name for path in Path(snapshot).iterdir()} + unexpected = sorted(actual - allowed) + if unexpected: + raise ValueError( + "KRONOS2's isolated runtime snapshot contains unexpected top-level " + f"entries: {unexpected}. Refusing to add it to the Python import path." + ) + + +@register("embedders", "kronos2") +class Kronos2Embedder(Embedder): + """KRONOS2 marker-aware ViT-B/16, returning the 768-dimensional CLS token.""" + + @property + def max_batch_size(self) -> int: + """Cap decoded batches before the CPU-staged, fixed-size forward path.""" + + return 8 + + def load( + self, + device: str = "cuda", + dtype: torch.dtype | None = None, # noqa: ARG002 - KRONOS2 is fixed fp32 + compile: bool = False, + ) -> Kronos2Embedder: + if compile: + raise ValueError( + "KRONOS2 does not support --compile in raw2features v0.2.1; " + "use the validated eager reference path" + ) + + import torch + + try: + from transformers import AutoModel + except ImportError as exc: # pragma: no cover - only without the extra + raise ImportError( + "KRONOS2 needs its pinned optional inference stack. Install " + "`raw2features[kronos2]`, log in to Hugging Face, and request " + "MahmoodLab/KRONOS2 access with an institutional account." + ) from exc + + snapshot = download_pinned_hf_snapshot( + self.spec.source, + self.spec.weights_revision, + allow_patterns=KRONOS2_SNAPSHOT_ALLOW_PATTERNS, + local_dir=pinned_model_cache_dir( + self.spec.source, self.spec.weights_revision + ), + ) + _assert_allowlisted_runtime_snapshot(snapshot) + weights = os.path.join(snapshot, str(self.spec.weights_filename)) + metadata = os.path.join(snapshot, KRONOS2_MARKER_METADATA_FILENAME) + # Both checks happen before Transformers imports or executes snapshot code. + verify_sha256(weights, self.spec.weights_sha256, what=self.spec.name) + verify_sha256( + metadata, + _metadata_sha256(self.spec), + what=f"{self.spec.name} marker metadata", + ) + + # Passing a local directory makes the pinned upstream override skip its own + # snapshot_download call. The outer Transformers resolver also receives + # local_files_only=True. Upstream temporarily inserts the directory at + # sys.path[0] for its bundled dinov2 package; remove that new entry after + # construction so its module names cannot shadow later imports. + snapshot_was_on_sys_path = snapshot in sys.path + try: + model = AutoModel.from_pretrained( + snapshot, + trust_remote_code=True, + local_files_only=True, + device=device, + ) + finally: + if not snapshot_was_on_sys_path: + sys.path[:] = [entry for entry in sys.path if entry != snapshot] + model = model.float().eval().to(device) + self._model = model + self._device = device + self._dtype = torch.float32 + self._metadata_path = metadata + self._load_marker_metadata(metadata) + self._panel: dict[str, Any] | None = None + # A successful load owns a fresh upstream model. Per-instance registration and + # panel-selection state from any prior unload/load cycle must not leak into it: + # a remembered digest without the corresponding appended model buffers would + # make the active vocabulary disagree with the actual forward path. + for attribute in ( + "_registered_additional_markers_digest", + "_novel_registration_failed", + "_raw2features_multiplex_source_indices", + "_binding_native_panel", + ): + self.__dict__.pop(attribute, None) + return self + + def _load_marker_metadata(self, path: str) -> None: + """Build the first-wins exact-key vocabulary from the verified CSV.""" + + index: dict[str, dict[str, Any]] = {} + compartments: set[str] = set() + families: set[str] = set() + with open(path, newline="", encoding="utf-8") as fh: + for row_number, row in enumerate(csv.DictReader(fh)): + canonical = str(row.get("marker_name") or "").strip() + if not canonical: + raise ValueError( + f"{self.spec.name}: {KRONOS2_MARKER_METADATA_FILENAME} " + f"contains a blank marker_name at data row {row_number + 1}" + ) + key = _marker_match_key(canonical) + record = { + "canonical": canonical, + "pretraining": str(row.get("pretraining") or "").strip().casefold() + == "yes", + "compartment": str(row.get("compartment") or ""), + "family": str(row.get("family") or ""), + } + compartments.add(record["compartment"]) + families.add(record["family"]) + # This mirrors upstream build_marker_index: CSV order wins when + # separator variants collapse to the same exact match key. + index.setdefault(key, record) + if not index: + raise ValueError( + f"{self.spec.name}: {KRONOS2_MARKER_METADATA_FILENAME} is empty" + ) + self._base_metadata_index = index + self._metadata_index = dict(index) + self._base_compartments = frozenset(compartments) + self._base_families = frozenset(families) + + def _published_metadata(self) -> dict[str, dict[str, Any]]: + """Return the verified shipped vocabulary, including in lightweight tests.""" + + return getattr(self, "_base_metadata_index", self._metadata_index) + + def _validated_novel_records( + self, contract: Any + ) -> tuple[list[dict[str, Any]], str, dict[str, dict[str, Any]]]: + """Validate novel rows against the verified base marker/category contract.""" + + rows, digest = _validate_additional_marker_contract(contract) + base_index = self._published_metadata() + allowed_compartments = getattr( + self, + "_base_compartments", + frozenset(record["compartment"] for record in base_index.values()), + ) + allowed_families = getattr( + self, + "_base_families", + frozenset(record["family"] for record in base_index.values()), + ) + additions: dict[str, dict[str, Any]] = {} + for row in rows: + canonical = str(row["marker_name"]) + match_key = _marker_match_key(canonical) + if match_key in base_index: + raise ValueError( + f"KRONOS2 additional marker {canonical!r} collides with the " + "verified published vocabulary under its exact matching contract" + ) + compartment = str(row["compartment"]) + if compartment not in allowed_compartments: + raise ValueError( + f"KRONOS2 additional marker {canonical!r} uses compartment " + f"{compartment!r}; it must exactly reuse one of the published " + f"categories: {sorted(allowed_compartments)}" + ) + family = str(row["family"]) + if family not in allowed_families: + raise ValueError( + f"KRONOS2 additional marker {canonical!r} uses family " + f"{family!r}; it must exactly reuse one of the published " + f"categories: {sorted(allowed_families)}" + ) + additions[match_key] = { + "canonical": canonical, + "pretraining": False, + "compartment": compartment, + "family": family, + } + return rows, digest, additions + + def _download_verified_biolinkbert_snapshot(self) -> str: + """Resolve and verify the complete conditional BioLinkBERT snapshot.""" + + snapshot = download_pinned_hf_snapshot( + KRONOS2_BIOLINKBERT_SOURCE, + KRONOS2_BIOLINKBERT_REVISION, + allow_patterns=tuple(KRONOS2_BIOLINKBERT_ARTIFACT_SHA256), + ) + for filename, sha256 in KRONOS2_BIOLINKBERT_ARTIFACT_SHA256.items(): + verify_sha256( + os.path.join(snapshot, filename), + sha256, + what=f"KRONOS2 BioLinkBERT {filename}", + ) + return snapshot + + def _register_novel_markers( + self, contract: dict[str, Any] + ) -> tuple[str, dict[str, dict[str, Any]]]: + """Run the pinned upstream registration path once for one semantic table.""" + + rows, digest, additions = self._validated_novel_records(contract) + if getattr(self, "_novel_registration_failed", False): + raise RuntimeError( + "this KRONOS2 instance had a failed novel-marker registration and " + "cannot be reused safely; unload it and construct a fresh model" + ) + registered_digest = getattr(self, "_registered_additional_markers_digest", None) + if registered_digest is not None: + if registered_digest != digest: + raise ValueError( + "this loaded KRONOS2 instance is already registered with " + f"additional-marker contract {registered_digest}; it cannot be " + f"reused with different contract {digest}. Load a fresh model." + ) + return digest, additions + + snapshot = self._download_verified_biolinkbert_snapshot() + marker_builder = self._model.backbone.marker_builder + builder_module_name = type(marker_builder).__module__ + try: + builder_module = importlib.import_module(builder_module_name) + encoder_cls = builder_module.BioLinkBERTTextEncoder + build_prompt = builder_module.build_prompt + except (ImportError, AttributeError) as exc: + raise RuntimeError( + "the pinned KRONOS2 marker-builder module does not expose its " + "BioLinkBERTTextEncoder and build_prompt registration contract" + ) from exc + + import torch + + with tempfile.TemporaryDirectory( + prefix="raw2features-kronos2-markers-" + ) as temp_dir: + work_dir = Path(temp_dir) + canonical_csv = work_dir / "additional_markers.csv" + materialize_kronos2_additional_markers_csv(contract, canonical_csv) + + cpu = torch.device("cpu") + encoder = encoder_cls( + model_id=snapshot, + device=cpu, + cache_dir=str(work_dir / "transformers-cache"), + max_length=128, + fp16_on_cuda=False, + ) + if not hasattr(encoder, "enc"): + raise RuntimeError( + "the pinned BioLinkBERTTextEncoder no longer exposes its encoder" + ) + encoder.enc.float().eval().to(cpu) + prompts = [ + build_prompt( + str(row["marker_name"]).lower(), + str(row["marker_full_name"]), + str(row["compartment"]), + str(row["family_desc"]), + ) + for row in rows + ] + with torch.inference_mode(): + encoded = encoder(prompts).detach().cpu().float() + expected_shape = (len(rows), 1024) + if tuple(encoded.shape) != expected_shape: + raise ValueError( + "KRONOS2 BioLinkBERT returned shape " + f"{tuple(encoded.shape)}, expected {expected_shape}" + ) + embeddings = np.asarray(encoded.numpy(), dtype=np.float32) + if not np.isfinite(embeddings).all(): + raise ValueError( + "KRONOS2 BioLinkBERT returned non-finite marker embeddings" + ) + np.save( + work_dir / "additional_marker_text_embeddings.npy", + np.ascontiguousarray(embeddings), + allow_pickle=False, + ) + del encoded, encoder + gc.collect() + + missing = object() + previous_cache = getattr(marker_builder, "cache_dir", missing) + marker_builder.cache_dir = work_dir + try: + self._model.register_additional_markers(str(canonical_csv)) + except Exception: + # Upstream updates the normalisation table before replacing all + # marker-builder buffers. A failure can therefore leave the live + # instance partially changed even though no store has been opened. + self._novel_registration_failed = True + raise + finally: + if previous_cache is missing: + delattr(marker_builder, "cache_dir") + else: + marker_builder.cache_dir = previous_cache + + self._registered_additional_markers_digest = digest + return digest, additions + + def bind_multiplex_panel( + self, + channel_names: list[str] | None, + *, + selected_channels: list[dict[str, Any]] | None = None, + model_params: dict[str, Any] | None = None, + ) -> dict: + """Bind a physical panel and conditionally register fingerprinted markers.""" + + if getattr(self, "_novel_registration_failed", False): + raise RuntimeError( + "this KRONOS2 instance had a failed novel-marker registration and " + "cannot be reused safely; unload it and construct a fresh model" + ) + params = {} if model_params is None else model_params + if not isinstance(params, dict) or set(params) - {"additional_markers"}: + keys = sorted(params) if isinstance(params, dict) else [] + raise ValueError( + "KRONOS2 native model parameters accept only " + f"'additional_markers'; received {keys or type(params).__name__}" + ) + + novel_summary: dict[str, Any] = { + "upstream_registration_available": True, + "enabled": False, + } + active_index = dict(self._published_metadata()) + if "additional_markers" in params: + contract = params["additional_markers"] + if not isinstance(contract, dict): + raise ValueError( + "KRONOS2 additional_markers must be a canonical parsed marker " + "contract, not a path, null value, or empty list" + ) + digest, additions = self._register_novel_markers(contract) + active_index.update(additions) + novel_summary = { + "upstream_registration_available": True, + "enabled": True, + "content_sha256": digest, + "markers": [str(row["marker_name"]) for row in contract["rows"]], + "registration": KRONOS2_NOVEL_MARKER_REGISTRATION_CONTRACT, + } + + # ``super`` validates and stores the selected physical source indices, then + # passes only those names to set_panel. Keep _panel indices local to that + # already-sliced tensor so the physical channels are never sliced twice. + self._metadata_index = active_index + self._binding_native_panel = True + try: + summary = super().bind_multiplex_panel( + channel_names, + selected_channels=selected_channels, + model_params=None, + ) + finally: + self._binding_native_panel = False + summary["novel_markers"] = novel_summary + return summary + + def set_panel(self, channel_names: list[str] | None) -> dict: + """Bind physical channels to published KRONOS2 markers before patch work. + + Matching is exact after the pinned separator-only normalisation. No KRONOSv1 + biological aliases or embedded-CD-token fallbacks are used. Named unknowns + fail now rather than receiving default statistics and crashing during forward. + """ + + if not channel_names: + raise ValueError( + "KRONOS2 requires one marker name for every multiplex channel" + ) + if not hasattr(self, "_metadata_index"): + raise RuntimeError("load KRONOS2 before binding a marker panel") + if not getattr(self, "_binding_native_panel", False): + # The compatibility API remains published-vocabulary-only. A prior + # explicit registration on a warm model must not make a later direct + # set_panel call silently accept a novel marker without its contract. + self._metadata_index = dict(self._published_metadata()) + + indices: list[int] = [] + marker_names: list[str] = [] + mapping: list[dict[str, Any]] = [] + dropped: list[str] = [] + unknown: list[str] = [] + non_pretraining: list[str] = [] + + for channel_index, value in enumerate(channel_names): + source_name = "" if value is None else str(value) + if not source_name.strip(): + dropped.append(source_name) + continue + record = self._metadata_index.get(_marker_match_key(source_name)) + if record is None: + unknown.append(source_name) + continue + canonical = str(record["canonical"]) + indices.append(channel_index) + marker_names.append(canonical) + if not record["pretraining"]: + non_pretraining.append(canonical) + mapping.append( + { + "channel": source_name, + "channel_index": channel_index, + "kronos2_marker": canonical, + "pretraining": bool(record["pretraining"]), + } + ) + + if unknown: + raise ValueError( + "KRONOS2 could not exactly match these named channels to its pinned " + f"marker_metadata.csv: {unknown}. No biological aliases are applied. " + "The published vocabulary includes DAPI and DRAQ5 but not common " + "Hoechst or DNA1/DNA2 labels. Such a channel can be omitted from an " + "explicit --marker selection while remaining available to nuclear " + "segmentation, or registered with dataset-specific statistics. " + "For intentionally novel markers, supply the complete additional-" + "marker CSV through KRONOS2's explicit additional-markers option; " + "raw2features will validate and fingerprint its statistics and pinned " + "BioLinkBERT registration before patch execution." + ) + if not indices: + raise ValueError( + "no non-empty channels matched the KRONOS2 marker vocabulary" + ) + + # Upstream's preferred_dapi argument forces DAPI statistics. Only the two + # published, explicitly reviewed nuclear stains are selected automatically; + # assigning an arbitrary novel ``dna_stain`` that override would change its + # user-supplied normalisation silently. + if "DAPI" in marker_names: + preferred_dapi = "DAPI" + elif "DRAQ5" in marker_names: + preferred_dapi = "DRAQ5" + else: + preferred_dapi = None + + self._panel = { + "idx": np.asarray(indices, dtype=np.int64), + "marker_names": marker_names, + "preferred_dapi": preferred_dapi, + "source_channel_count": len(channel_names), + } + return { + "n_markers": len(marker_names), + "kept": [str(channel_names[i]) for i in indices], + "dropped": dropped, + "unmatched": [], + "matched": marker_names, + "defaulted": [], + "non_pretraining": non_pretraining, + "preferred_dapi": preferred_dapi, + "mapping": mapping, + "matching": KRONOS2_MARKER_MATCHING_CONTRACT, + "vocabulary": { + "source": self.spec.source, + "filename": KRONOS2_MARKER_METADATA_FILENAME, + "revision": self.spec.weights_revision, + "sha256": _metadata_sha256(self.spec), + }, + "novel_markers": { + "upstream_registration_available": True, + "enabled": False, + }, + } + + def transform_batch( + self, + patches_hwc: list[np.ndarray], + device: str, # noqa: ARG002 - KRONOS2 stages transforms on CPU + ) -> torch.Tensor: + """Scale raw HWC marker patches, then delegate z-scoring upstream.""" + + if self._panel is None: + raise RuntimeError("call set_panel() before embedding multiplex patches") + if not patches_hwc: + raise ValueError("KRONOS2 transform_batch received an empty patch batch") + expected_channels = int(self._panel["source_channel_count"]) + source_dtype: np.dtype | None = None + side_px: int | None = None + validated: list[np.ndarray] = [] + for patch_index, patch in enumerate(patches_hwc): + array = np.asarray(patch) + if array.ndim != 3: + raise ValueError( + f"KRONOS2 patch {patch_index} must be HWC, got shape " + f"{tuple(array.shape)}" + ) + height, width, channels = array.shape + if height <= 0 or height != width or height % 16: + raise ValueError( + f"KRONOS2 patch {patch_index} must be square with a positive " + f"side divisible by its 16-pixel stride, got {height}x{width}" + ) + if channels != expected_channels: + raise ValueError( + f"KRONOS2 patch {patch_index} has {channels} channels but the " + f"bound selected panel has {expected_channels}" + ) + if side_px is not None and height != side_px: + raise ValueError( + "KRONOS2 requires one spatial patch size across a batch" + ) + dtype = np.dtype(array.dtype) + if source_dtype is not None and dtype != source_dtype: + raise ValueError( + "KRONOS2 requires one source dtype across a patch batch" + ) + if np.issubdtype(dtype, np.floating) and not np.isfinite(array).all(): + raise ValueError( + f"KRONOS2 patch {patch_index} contains non-finite float values" + ) + source_dtype = dtype + side_px = height + validated.append(array) + if source_dtype is None: # pragma: no cover - guarded by non-empty input + raise AssertionError("KRONOS2 source dtype was not resolved") + selected = np.stack( + [ + patch[:, :, self._panel["idx"]].astype(np.float32, copy=False) + for patch in validated + ] + ) + selected = np.transpose(selected, (0, 3, 1, 2)) + scaled = selected / _scaling_factor(source_dtype) + normalized = self._model.preprocess( + scaled, + self._panel["marker_names"], + preferred_dapi=self._panel["preferred_dapi"], + ) + import torch + + return torch.from_numpy( + np.ascontiguousarray(np.asarray(normalized, dtype=np.float32)) + ) + + def embed_batch(self, batch: torch.Tensor) -> torch.Tensor: + """Return CLS vectors through the upstream's fixed eight-row regime.""" + + import torch + + if self._panel is None: + raise RuntimeError("call set_panel() before embedding multiplex patches") + if batch.ndim != 4 or batch.shape[0] <= 0: + raise ValueError("KRONOS2 embed_batch expects a non-empty BCHW tensor") + if batch.shape[1] != len(self._panel["marker_names"]): + raise ValueError( + f"KRONOS2 batch has {batch.shape[1]} marker channels but the bound " + f"panel has {len(self._panel['marker_names'])}" + ) + outputs: list[torch.Tensor] = [] + with torch.inference_mode(): + for start in range(0, int(batch.shape[0]), 8): + microbatch = batch[start : start + 8] + valid_rows = int(microbatch.shape[0]) + if valid_rows < 8: + repeats = microbatch[-1:].expand( + 8 - valid_rows, *microbatch.shape[1:] + ) + microbatch = torch.cat((microbatch, repeats), dim=0) + cls = self._model( + microbatch.to(self._device, dtype=torch.float32), + self._panel["marker_names"], + ) + if tuple(cls.shape) != (8, self.embedding_dim): + raise ValueError( + "KRONOS2 returned shape " + f"{tuple(cls.shape)}, expected {(8, self.embedding_dim)}" + ) + outputs.append(cls[:valid_rows].float().cpu()) + return torch.cat(outputs, dim=0) diff --git a/src/raw2features/embedders/kronos2_metadata.py b/src/raw2features/embedders/kronos2_metadata.py new file mode 100644 index 0000000..9dec50d --- /dev/null +++ b/src/raw2features/embedders/kronos2_metadata.py @@ -0,0 +1,362 @@ +"""Canonical input contract for user-supplied KRONOS2 marker metadata. + +The upstream KRONOS2 novel-marker path consumes a small CSV table. This module +normalises only the fields that path reads, validates the parts that can be +checked without the gated base vocabulary, and hashes the semantic table together +with the immutable text-registration recipe rather than incidental CSV formatting. +""" + +from __future__ import annotations + +import csv +import hashlib +import io +import json +import math +import os +import unicodedata +from collections.abc import Mapping +from copy import deepcopy +from typing import Any, TextIO + +KRONOS2_ADDITIONAL_MARKER_CONTRACT_VERSION = 1 +KRONOS2_ADDITIONAL_MARKER_COLUMNS = ( + "marker_name", + "marker_full_name", + "compartment", + "family", + "family_desc", + "mean", + "std", +) +_TEXT_COLUMNS = KRONOS2_ADDITIONAL_MARKER_COLUMNS[:5] + +KRONOS2_BIOLINKBERT_SOURCE = "hf_hub:michiyasunaga/BioLinkBERT-large" +KRONOS2_BIOLINKBERT_REVISION = "1eb6d81c5fc1c42d3a43c71956b0e526558ae053" +KRONOS2_BIOLINKBERT_ARTIFACT_SHA256 = { + "pytorch_model.bin": ( + "fed75e5716547b54198d4dd123e7a3f3c64a82e1172b3492a11deebd6ab4cd4d" + ), + "config.json": ("ba7be72cf4013c5a69166e88d51f922aee44bfde6927cf9315580c23a6e020a6"), + "tokenizer.json": ( + "5a797027c356bfdc779fd021b4a7a2b2f341242ddd7035d385f56b8e242dac2a" + ), + "tokenizer_config.json": ( + "98632482a9851173e796e4366ede775ddc2900564a0ac659acb89248e39fbbd3" + ), + "special_tokens_map.json": ( + "303df45a03609e4ead04bc3dc1536d0ab19b5358db685b6f3da123d05ec200e3" + ), + "vocab.txt": ("7b36651908a88bc38bda41b728b2a598191e0d3b553cbacf7b1e5f026d5b5b9f"), +} + +# This conditional contract is included only when a user supplies an +# additional-marker table. Published-marker inference therefore retains its smaller +# fingerprint and never needs BioLinkBERT, while a novel-marker result binds every +# part of the upstream text-embedding recipe. +KRONOS2_NOVEL_MARKER_REGISTRATION_CONTRACT: dict[str, Any] = { + "contract_version": 1, + "activation": "additional_markers_parameter_only", + "registration": { + "entrypoint": "pinned_kronos2_model.register_additional_markers", + "canonical_csv_columns": list(KRONOS2_ADDITIONAL_MARKER_COLUMNS), + "statistics": "user_supplied_per_marker_mean_std", + "category_policy": "reuse_exact_base_compartment_and_family_only", + "base_marker_collision_policy": "reject_separator_insensitive_exact", + "text_embedding_handoff": "temporary_additional_marker_text_embeddings.npy", + }, + "prompt": { + "builder": "pinned_kronos2_builder_module.build_prompt", + "inputs": [ + "marker_name_lowercase", + "marker_full_name", + "compartment", + "family_desc", + ], + "template": ( + "Marker: {marker_name.upper()} ({marker_full_name}) | " + "It is located in the {compartment} compartment and " + "belongs to the {family_desc} family" + ), + }, + "text_encoder": { + "class": "pinned_kronos2_builder_module.BioLinkBERTTextEncoder", + "source": KRONOS2_BIOLINKBERT_SOURCE, + "revision": KRONOS2_BIOLINKBERT_REVISION, + "license": "Apache-2.0", + "artifact_sha256": deepcopy(KRONOS2_BIOLINKBERT_ARTIFACT_SHA256), + "input": "verified_pinned_local_snapshot", + "device": "cpu", + "precision": "fp32", + "embedding_dim": 1024, + "tokenizer": { + "use_fast": True, + "padding": True, + "truncation": True, + "max_length": 128, + }, + "pooling": "attention_mask_mean_last_hidden_state_including_special_tokens", + "normalization": "l2", + "fp16_on_cuda": False, + }, +} + + +class Kronos2MarkerMetadataError(ValueError): + """Raised when an additional-marker CSV cannot form a safe contract.""" + + +def _canonical_json(value: Any) -> str: + return json.dumps( + value, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + allow_nan=False, + ) + + +def _contains_control(value: str) -> bool: + return any(unicodedata.category(character) == "Cc" for character in value) + + +def _canonical_text(value: str | None, *, column: str, row_number: int) -> str: + if value is None: + value = "" + # Prompt text is content, not CSV formatting. Preserve its Unicode after + # trimming so BioLinkBERT receives what the user supplied and the semantic digest + # moves whenever that text changes. + value = str(value) + if _contains_control(value): + raise Kronos2MarkerMetadataError( + f"row {row_number} column {column!r} contains a control character" + ) + value = value.strip() + if not value: + raise Kronos2MarkerMetadataError( + f"row {row_number} column {column!r} must not be blank" + ) + return value + + +def _canonical_number(value: str | None, *, column: str, row_number: int) -> float: + if value is None: + value = "" + value = str(value) + if _contains_control(value): + raise Kronos2MarkerMetadataError( + f"row {row_number} column {column!r} contains a control character" + ) + try: + number = float(value.strip()) + except ValueError as exc: + raise Kronos2MarkerMetadataError( + f"row {row_number} column {column!r} must be a finite number" + ) from exc + if not math.isfinite(number): + raise Kronos2MarkerMetadataError( + f"row {row_number} column {column!r} must be a finite number" + ) + if column == "std" and number <= 0: + raise Kronos2MarkerMetadataError( + f"row {row_number} column 'std' must be greater than zero" + ) + # Treat signed zero as one canonical numeric value. + return 0.0 if number == 0 else number + + +def _kronos2_marker_match_key(value: str) -> str: + """Mirror KRONOS2's separator-insensitive forward-lookup identity.""" + + cleaned = ( + value.lower() + .translate( + str.maketrans( + { + "-": "_", + " ": "_", + ":": "_", + "α": "a", + "(": "_", + ")": "_", + "/": "", + } + ) + ) + .strip() + ) + return cleaned.replace("_", "").replace(".", "") + + +def _validate_headers(fieldnames: list[str | None] | None) -> None: + if not fieldnames: + raise Kronos2MarkerMetadataError("additional-marker CSV has no header") + if any(name is None or not name for name in fieldnames): + raise Kronos2MarkerMetadataError( + "additional-marker CSV contains a blank column name" + ) + duplicate_headers = sorted( + {name for name in fieldnames if fieldnames.count(name) > 1} + ) + if duplicate_headers: + raise Kronos2MarkerMetadataError( + f"additional-marker CSV has duplicate columns: {duplicate_headers}" + ) + missing = [ + column + for column in KRONOS2_ADDITIONAL_MARKER_COLUMNS + if column not in fieldnames + ] + if missing: + raise Kronos2MarkerMetadataError( + f"additional-marker CSV is missing required columns: {missing}" + ) + + +def _canonical_payload(rows: list[dict[str, Any]]) -> dict[str, Any]: + return { + "contract_version": KRONOS2_ADDITIONAL_MARKER_CONTRACT_VERSION, + "consumed_columns": list(KRONOS2_ADDITIONAL_MARKER_COLUMNS), + "registration": deepcopy(KRONOS2_NOVEL_MARKER_REGISTRATION_CONTRACT), + "rows": rows, + } + + +def parse_kronos2_additional_markers( + path: str | os.PathLike[str], + *, + include_original_file_sha256: bool = False, +) -> dict[str, Any]: + """Parse *path* into a JSON-safe, fingerprintable marker-table contract. + + Extra CSV columns are deliberately ignored: the contract lists and hashes + exactly the columns consumed by KRONOS2. Row order is semantic because the + upstream registration path appends marker buffers in that order. Validation + against the gated base vocabulary's compartment/family categories is left to + the loader after that independently verified metadata is available. The + original file-byte digest is optional audit data and is omitted by default so + formatting-only CSV changes do not change an embedding's output identity. + """ + + try: + raw = os.fsdecode(path) + with open(raw, "rb") as handle: + file_bytes = handle.read() + except OSError as exc: + raise Kronos2MarkerMetadataError( + f"could not read additional-marker CSV {os.fspath(path)!r}: {exc}" + ) from exc + + try: + text = file_bytes.decode("utf-8-sig") + except UnicodeDecodeError as exc: + raise Kronos2MarkerMetadataError( + "additional-marker CSV must be UTF-8 encoded" + ) from exc + + rows: list[dict[str, Any]] = [] + try: + reader = csv.DictReader(io.StringIO(text, newline=""), strict=True) + _validate_headers(reader.fieldnames) + for row_number, source in enumerate(reader, start=2): + if None in source: + raise Kronos2MarkerMetadataError( + f"row {row_number} contains more fields than the CSV header" + ) + row = { + column: _canonical_text( + source.get(column), column=column, row_number=row_number + ) + for column in _TEXT_COLUMNS + } + row["mean"] = _canonical_number( + source.get("mean"), column="mean", row_number=row_number + ) + row["std"] = _canonical_number( + source.get("std"), column="std", row_number=row_number + ) + rows.append(row) + except csv.Error as exc: + raise Kronos2MarkerMetadataError( + f"could not parse additional-marker CSV: {exc}" + ) from exc + + if not rows: + raise Kronos2MarkerMetadataError( + "additional-marker CSV must contain at least one marker row" + ) + + casefold_seen: dict[str, tuple[int, str]] = {} + match_seen: dict[str, tuple[int, str]] = {} + for row_number, row in enumerate(rows, start=2): + marker_name = row["marker_name"] + casefold_key = marker_name.casefold() + match_key = _kronos2_marker_match_key(marker_name) + if previous := casefold_seen.get(casefold_key): + raise Kronos2MarkerMetadataError( + "marker_name values must be unique case-insensitively: " + f"rows {previous[0]} ({previous[1]!r}) and " + f"{row_number} ({marker_name!r})" + ) + if previous := match_seen.get(match_key): + raise Kronos2MarkerMetadataError( + "marker_name values must be unique under KRONOS2's " + "separator-insensitive lookup: " + f"rows {previous[0]} ({previous[1]!r}) and " + f"{row_number} ({marker_name!r})" + ) + casefold_seen[casefold_key] = (row_number, marker_name) + match_seen[match_key] = (row_number, marker_name) + + semantic_payload = _canonical_payload(rows) + payload = deepcopy(semantic_payload) + payload["content_sha256"] = hashlib.sha256( + _canonical_json(semantic_payload).encode("utf-8") + ).hexdigest() + if include_original_file_sha256: + payload["original_file_sha256"] = hashlib.sha256(file_bytes).hexdigest() + return payload + + +def materialize_kronos2_additional_markers_csv( + contract: Mapping[str, Any], + destination: str | os.PathLike[str] | TextIO, +) -> None: + """Write a parsed contract as the exact canonical CSV consumed upstream.""" + + if contract.get("contract_version") != KRONOS2_ADDITIONAL_MARKER_CONTRACT_VERSION: + raise Kronos2MarkerMetadataError( + "cannot materialize an unknown additional-marker contract version" + ) + if contract.get("consumed_columns") != list(KRONOS2_ADDITIONAL_MARKER_COLUMNS): + raise Kronos2MarkerMetadataError( + "cannot materialize a contract with different consumed columns" + ) + if contract.get("registration") != KRONOS2_NOVEL_MARKER_REGISTRATION_CONTRACT: + raise Kronos2MarkerMetadataError( + "cannot materialize a contract with a different registration recipe" + ) + rows = contract.get("rows") + if not isinstance(rows, list) or not rows: + raise Kronos2MarkerMetadataError( + "cannot materialize a contract without marker rows" + ) + + close = False + if hasattr(destination, "write"): + handle = destination + else: + handle = open(destination, "w", encoding="utf-8", newline="") + close = True + try: + writer = csv.DictWriter( + handle, + fieldnames=KRONOS2_ADDITIONAL_MARKER_COLUMNS, + extrasaction="ignore", + lineterminator="\n", + ) + writer.writeheader() + writer.writerows(rows) + finally: + if close: + handle.close() diff --git a/src/raw2features/embedders/registry.yaml b/src/raw2features/embedders/registry.yaml index f6517c8..08c2b20 100644 --- a/src/raw2features/embedders/registry.yaml +++ b/src/raw2features/embedders/registry.yaml @@ -1011,6 +1011,42 @@ kronos: SDPA (flash attention) by default -- no xFormers; much faster and far lighter on memory than KRONOS's upstream naive-attention fallback, so the default batch fits. Opt out with `sdpa: false`. Validated full-slide on real CODEX multiplex data. +kronos2: + family: kronos2 + modality: multiplex + source: "hf_hub:MahmoodLab/KRONOS2" + embedding_dim: 768 + input_size: 256 + recommended_patch_px: 256 + pooling: cls + # RGB normalization is unused. Raw native marker channels first use the + # upstream dtype divisor, then the verified marker_metadata.csv statistics. + mean: [0.0, 0.0, 0.0] + std: [1.0, 1.0, 1.0] + interpolation: bicubic + inference_amp: fp32 + transform_source_url: "https://huggingface.co/MahmoodLab/KRONOS2" + weights_filename: "kronos2_vitb16_teacher.pth" + weights_sha256: "afa261ef8b63bb0acea62612d7594ac11fbdb89fbdd23c513f046a034725d2b4" + weights_revision: "e2a5473206c02d9973885b6445ef783f7e082951" + doi: "10.48550/arXiv.2506.03373" + license: "CC-BY-NC-ND-4.0 (KRONOS2; gated non-commercial academic use)" + gated: true + timm_kwargs: + marker_metadata_filename: "marker_metadata.csv" + marker_metadata_sha256: "1e6791448a792c0edec9b2ba74efceaefba84d96852351cdc58d67f80b539074" + notes: >- + KRONOS2 (MahmoodLab) is a native marker-aware multiplex ViT-B/16 that returns + a 768-d CLS vector. The complete pinned Hugging Face snapshot is loaded locally + after the weight and marker-metadata hashes are verified. Marker names match + exactly after the upstream separator-only normalization; no biological aliases + or KRONOSv1 compound-token fallback is applied. Published marker text embeddings + are checkpoint buffers, so ordinary inference does not download BioLinkBERT. + An explicit complete additional-marker CSV enables the upstream registration path + with user-supplied statistics and a separately pinned, SHA-verified BioLinkBERT + snapshot; that conditional contract is fingerprinted. Requires the `kronos2` + extra plus individually approved gated access. + musk: family: musk source: "hf-hub:xiangjx/musk" diff --git a/src/raw2features/multiplex/__init__.py b/src/raw2features/multiplex/__init__.py index 1397d65..d2dd105 100644 --- a/src/raw2features/multiplex/__init__.py +++ b/src/raw2features/multiplex/__init__.py @@ -6,6 +6,7 @@ PreparedMultiplexStrategy, ) from .channelwise import ChannelwiseStrategy +from .panel import marker_name_identity, resolve_marker_selection from .registry import build_strategy __all__ = [ @@ -14,4 +15,6 @@ "MultiplexStrategy", "PreparedMultiplexStrategy", "build_strategy", + "marker_name_identity", + "resolve_marker_selection", ] diff --git a/src/raw2features/multiplex/panel.py b/src/raw2features/multiplex/panel.py new file mode 100644 index 0000000..a4ce45b --- /dev/null +++ b/src/raw2features/multiplex/panel.py @@ -0,0 +1,73 @@ +"""Deterministic selection of named channels from a positional multiplex panel.""" + +from __future__ import annotations + +import unicodedata +from collections.abc import Sequence +from typing import Any + + +def marker_name_identity(value: Any) -> str: + """Return the user-facing marker identity used by ``--marker`` selection. + + This is deliberately less biological than a model vocabulary resolver. It only + makes metadata spelling comparisons Unicode-, whitespace-, and case-insensitive; + each native model remains responsible for mapping the selected source labels to + its own canonical marker identities. + """ + + return unicodedata.normalize("NFKC", str(value)).strip().casefold() + + +def resolve_marker_selection( + channel_names: Sequence[str | None], + requested_markers: Sequence[str] | None = None, +) -> list[dict[str, Any]]: + """Resolve an ordered marker request to physical C-axis positions. + + With no explicit request every physical position is retained, including unnamed + slots. Keeping those slots is important because model-specific panel binding may + drop them later, but their positions still define the source array. An explicit + request must identify each channel uniquely after :func:`marker_name_identity`. + The returned order is the requested order, not source order. + """ + + panel = ["" if value is None else str(value) for value in channel_names] + requested = [str(value).strip() for value in (requested_markers or [])] + if any(not value for value in requested): + raise ValueError("selected multiplex marker names must be non-empty") + + if not requested: + return [ + {"source_index": index, "source_name": source_name} + for index, source_name in enumerate(panel) + ] + + positions: dict[str, list[int]] = {} + for index, source_name in enumerate(panel): + identity = marker_name_identity(source_name) + if identity: + positions.setdefault(identity, []).append(index) + + records: list[dict[str, Any]] = [] + seen: set[str] = set() + for requested_name in requested: + identity = marker_name_identity(requested_name) + if identity in seen: + raise ValueError( + f"multiplex marker {requested_name!r} was requested more than once" + ) + seen.add(identity) + matches = positions.get(identity, []) + if not matches: + raise ValueError( + f"requested multiplex marker {requested_name!r} is not present" + ) + if len(matches) != 1: + raise ValueError( + f"requested multiplex marker {requested_name!r} is ambiguous; " + f"it matches source channel indices {matches}" + ) + index = matches[0] + records.append({"source_index": index, "source_name": panel[index]}) + return records diff --git a/src/raw2features/pipeline/runner.py b/src/raw2features/pipeline/runner.py index 600d635..45b94e5 100644 --- a/src/raw2features/pipeline/runner.py +++ b/src/raw2features/pipeline/runner.py @@ -106,6 +106,77 @@ def _models_need_nuclear(models: list[str], specs: dict | None = None) -> bool: return _registered_models_need_nuclear(missing) +def _native_multiplex_model_names( + models: list[str], specs: dict | None = None +) -> list[str]: + """Return requested native multiplex model names in request order.""" + + provided = specs or {} + native: list[str] = [] + for name in models: + spec = provided.get(name) + if spec is None: + try: + spec = get_spec(name) + except KeyError: + continue + if getattr(spec, "modality", "brightfield") == "multiplex": + native.append(name) + return native + + +def _validate_native_patch_geometry( + cfg: RunConfig, specs: dict | None = None +) -> None: + """Reject a native model's unsupported patch geometry before store mutation.""" + + provided = specs or {} + for name in cfg.models: + spec = provided.get(name) + if spec is None: + try: + spec = get_spec(name) + except KeyError: + continue + if getattr(spec, "family", "") != "kronos2": + continue + from raw2features.embedders.kronos2_embedder import KRONOS2_FORWARD_CONTRACT + + stride = int(KRONOS2_FORWARD_CONTRACT["upstream_patch_stride_px"]) + if cfg.patch_px % stride: + raise ValueError( + f"KRONOS2 patch size must be divisible by {stride}; " + f"received {cfg.patch_px}" + ) + + +def _reject_mixed_native_and_brightfield_models( + models: list[str], specs: dict | None = None +) -> None: + """Reject one request that would feed the same decoded patches to both modes.""" + + provided = specs or {} + modalities: dict[str, str] = {} + for name in models: + spec = provided.get(name) + if spec is None: + try: + spec = get_spec(name) + except KeyError: + continue + modalities[name] = getattr(spec, "modality", "brightfield") + native = [name for name, modality in modalities.items() if modality == "multiplex"] + brightfield = [ + name for name, modality in modalities.items() if modality != "multiplex" + ] + if native and brightfield: + raise ValueError( + "native multiplex and brightfield patch encoders cannot be combined in " + f"one request (native={native}, brightfield={brightfield}). Run them " + "separately so each model receives the correct source-channel view." + ) + + @dataclass class RunConfig: """Everything needed to embed one slide. Content-affecting fields feed the @@ -200,6 +271,14 @@ class RunConfig: # runtime migration evidence only: effective names above determine current output, # while the original list lets us prove whether a v0.1 selector used the same index. resolved_original_channel_names: list[str] = field(default_factory=list) + # Ordered physical source channels selected for native multiplex embedders. Empty + # until source resolution; each record is {source_index, source_name}. The full + # resolved panel remains available for segmentation and provenance. + resolved_native_marker_selection: list[dict[str, Any]] = field(default_factory=list) + # Namespaced, finite JSON configuration for native multiplex model adapters. + # Example: {"kronos2": {...}}. No built-in option is silently accepted: each + # embedder's bind_multiplex_panel hook validates its own namespace. + native_multiplex_params: dict[str, dict[str, Any]] = field(default_factory=dict) def __post_init__(self) -> None: """Reject invalid geometry before hashing, model loading, or store writes.""" @@ -210,6 +289,32 @@ def __post_init__(self) -> None: self.step_px = _positive_int(self.step_px, field="step_px") if self.source_mpp is not None: self.source_mpp = _positive_float(self.source_mpp, field="source_mpp") + if self.amp in {"bf16", "fp16"}: + fixed_fp32 = [] + for model_name in self.models: + try: + if get_spec(model_name).family == "kronos2": + fixed_fp32.append(model_name) + except KeyError: + continue + if fixed_fp32: + raise ValueError( + "KRONOS2 follows its published fp32 inference path and does not " + f"accept --amp {self.amp}; use --amp auto or --amp fp32" + ) + if self.compile: + unvalidated_compile = [] + for model_name in self.models: + try: + if get_spec(model_name).family == "kronos2": + unvalidated_compile.append(model_name) + except KeyError: + continue + if unvalidated_compile: + raise ValueError( + "KRONOS2 does not support --compile in raw2features v0.2.1; " + "use the validated eager reference path" + ) if self.multiplex_strategy is not None: self.multiplex_strategy = self.multiplex_strategy.strip().lower() or None if not self.no_seg and ( @@ -221,6 +326,8 @@ def __post_init__(self) -> None: # Otsu brightfield grid whose coordinates differ. self.segmenter = "nuclear" self.multiplex_markers = [str(name).strip() for name in self.multiplex_markers] + if any(not name for name in self.multiplex_markers): + raise ValueError("multiplex marker names must be non-empty") self.channel_names_override = [ str(name).strip() for name in self.channel_names_override ] @@ -233,6 +340,49 @@ def __post_init__(self) -> None: "" if name is None else str(name) for name in self.resolved_original_channel_names ] + if not isinstance(self.resolved_native_marker_selection, list): + raise ValueError("resolved_native_marker_selection must be a list") + resolved_selection: list[dict[str, Any]] = [] + seen_selected_indices: set[int] = set() + for record in self.resolved_native_marker_selection: + if not isinstance(record, dict) or set(record) != { + "source_index", + "source_name", + }: + raise ValueError( + "resolved_native_marker_selection entries must contain exactly " + "source_index and source_name" + ) + index = record["source_index"] + name = record["source_name"] + if ( + isinstance(index, bool) + or not isinstance(index, int) + or index < 0 + or not isinstance(name, str) + ): + raise ValueError( + "resolved_native_marker_selection entries require a non-negative " + "integer source_index and string source_name" + ) + if index in seen_selected_indices: + raise ValueError( + "resolved_native_marker_selection source indices must be unique" + ) + seen_selected_indices.add(index) + resolved_selection.append({"source_index": index, "source_name": name}) + self.resolved_native_marker_selection = resolved_selection + if self.resolved_channel_names: + for record in self.resolved_native_marker_selection: + index = record["source_index"] + if ( + index >= len(self.resolved_channel_names) + or self.resolved_channel_names[index] != record["source_name"] + ): + raise ValueError( + "resolved_native_marker_selection does not match " + "resolved_channel_names" + ) self.resolved_nuclear_channel_indices = list( self.resolved_nuclear_channel_indices ) @@ -265,10 +415,28 @@ def __post_init__(self) -> None: raise ValueError( "multiplex_strategy_params must contain only finite JSON values" ) from exc + if not isinstance(self.native_multiplex_params, dict) or any( + not isinstance(key, str) or not isinstance(value, dict) + for key, value in self.native_multiplex_params.items() + ): + raise ValueError( + "native_multiplex_params must be a mapping from model names to " + "parameter mappings" + ) + try: + encoded_native_params = json.dumps( + self.native_multiplex_params, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ) + self.native_multiplex_params = json.loads(encoded_native_params) + except (TypeError, ValueError) as exc: + raise ValueError( + "native_multiplex_params must contain only finite JSON values" + ) from exc if self.multiplex_strategy is None: - orphaned = bool( - self.multiplex_markers or self.multiplex_strategy_params - ) or any( + strategy_orphaned = bool(self.multiplex_strategy_params) or any( ( str(self.multiplex_normalization).lower() != "percentile", self.multiplex_percentile_low != 1.0, @@ -277,17 +445,35 @@ def __post_init__(self) -> None: self.multiplex_normalization_max_side_px != 2048, ) ) - if orphaned: + if strategy_orphaned: raise ValueError( - "multiplex marker/normalization/aggregation options require " + "multiplex normalization/aggregation/strategy options require " "multiplex_strategy (for example 'channelwise')" ) + # Preserve the ordinary-model typo guard while allowing an unknown name + # to be supplied later as an injected native ModelSpec. The definitive + # check happens in resolve_multiplex_source_config, where those external + # specifications are available. + if self.multiplex_markers or self.native_multiplex_params: + registered_modalities: list[str] = [] + has_unknown = False + for model_name in self.models: + try: + registered_modalities.append(get_spec(model_name).modality) + except KeyError: + has_unknown = True + if ( + not has_unknown + and "multiplex" not in registered_modalities + ): + raise ValueError( + "multiplex marker/native-model options require " + "multiplex_strategy or a native multiplex model" + ) # Ordinary brightfield requests deliberately do not validate or hash the # dormant strategy defaults. Their historical identity and execution path # remain byte-for-byte unchanged. return - if any(not name for name in self.multiplex_markers): - raise ValueError("multiplex marker names must be non-empty") self.multiplex_normalization = str(self.multiplex_normalization).lower() if self.multiplex_normalization not in {"percentile", "dtype"}: raise ValueError( @@ -334,6 +520,8 @@ def __post_init__(self) -> None: "multiplex_strategy_params", "multiplex_normalization_max_side_px", "resolved_channel_names", + "resolved_native_marker_selection", + "native_multiplex_params", "resolved_nuclear_channel_indices", "snap_to_level", "mpp_tolerance", @@ -366,6 +554,7 @@ def _hash_payload( include_models: bool, include_amp: bool = True, models: list[str] | None = None, + amp_override: str | None = None, segmenter_override: str | None = None, ) -> str: # Build the hashed payload from _CONTENT_FIELDS, applying the few value @@ -373,6 +562,27 @@ def _hash_payload( # irrelevant -- the bytes are identical to listing the dict by hand. payload: dict = {} for name in self._CONTENT_FIELDS: + if name == "resolved_native_marker_selection": + # The all-channel default is already represented by the complete + # resolved panel and must not perturb historical KRONOSv1 identities. + if ( + not include_models + or not self.multiplex_markers + or not self.resolved_native_marker_selection + ): + continue + if name == "native_multiplex_params": + if not include_models or not self.native_multiplex_params: + continue + selected_models = set(models if models is not None else self.models) + value = { + model: params + for model, params in self.native_multiplex_params.items() + if model in selected_models + } + if value: + payload[name] = value + continue if name == "resolved_channel_names": # The complete panel changes model outputs/receipts, but only the # resolved nuclear position below is a geometry input. @@ -414,7 +624,12 @@ def _hash_payload( payload[name] = sorted(models if models is not None else self.models) continue value = getattr(self, name) - if name == "segmenter": + if name == "amp" and amp_override is not None: + # Legacy-grid discovery enumerates old AMP-bearing identities even + # for current loaders that accept only fp32. This is a migration hash + # input, not a runnable request, so do not construct a RunConfig for it. + value = amp_override + elif name == "segmenter": value = ( "none" if self.no_seg @@ -470,17 +685,20 @@ def compatible_legacy_grid_hashes(self) -> tuple[str, ...]: existing compatibility families. """ + legacy_base = replace( + self, + resolved_channel_names=[], + resolved_nuclear_channel_indices=[], + ) amps = (self.amp, "auto", "fp32", "bf16", "fp16") - configs = [ - replace( - self, - amp=value, - resolved_channel_names=[], - resolved_nuclear_channel_indices=[], + hashes = [ + legacy_base._hash_payload( + include_models=False, + include_amp=True, + amp_override=value, ) for value in amps ] - hashes = [config.legacy_grid_hash() for config in configs] if self._is_native_multiplex_grid_request() and not ( self._v01_nuclear_selector_matches_current() ): @@ -509,13 +727,17 @@ def compatible_legacy_grid_segmenters(self) -> dict[str, str]: self._v01_nuclear_selector_matches_current() ): return {} + legacy_base = replace( + self, + resolved_channel_names=[], + resolved_nuclear_channel_indices=[], + ) nuclear_hashes = [ - replace( - self, - amp=value, - resolved_channel_names=[], - resolved_nuclear_channel_indices=[], - ).legacy_grid_hash() + legacy_base._hash_payload( + include_models=False, + include_amp=True, + amp_override=value, + ) for value in (self.amp, "auto", "fp32", "bf16", "fp16") ] return { @@ -562,15 +784,6 @@ def _native_multiplex_otsu_grid_aliases(self) -> tuple[str, ...]: ): return () amps = (self.amp, "auto", "fp32", "bf16", "fp16") - configs = [ - replace( - self, - amp=value, - resolved_channel_names=[], - resolved_nuclear_channel_indices=[], - ) - for value in amps - ] legacy_base = replace( self, resolved_channel_names=[], @@ -584,12 +797,13 @@ def _native_multiplex_otsu_grid_aliases(self) -> tuple[str, ...]: ) ] hashes.extend( - config._hash_payload( + legacy_base._hash_payload( include_models=False, include_amp=True, + amp_override=value, segmenter_override="otsu", ) - for config in configs + for value in amps ) return tuple(dict.fromkeys(hashes)) @@ -681,7 +895,26 @@ def resolve_multiplex_source_config( absent or blank OME metadata, but the reader rejects conflicts with existing labels. """ - native_multiplex = _models_need_nuclear(cfg.models, model_specs) + if cfg.multiplex_strategy is None: + _reject_mixed_native_and_brightfield_models(cfg.models, model_specs) + native_models = _native_multiplex_model_names(cfg.models, model_specs) + native_multiplex = bool(native_models) + unknown_param_namespaces = sorted( + set(cfg.native_multiplex_params) - set(native_models) + ) + if unknown_param_namespaces: + raise ValueError( + "native multiplex parameters were supplied for models that are not " + f"requested native multiplex encoders: {unknown_param_namespaces}" + ) + if ( + cfg.multiplex_markers + and cfg.multiplex_strategy is None + and not native_multiplex + ): + raise ValueError( + "--marker requires a native multiplex model or --multiplex-strategy" + ) if ( cfg.channel_names_override and cfg.multiplex_strategy is None @@ -730,6 +963,18 @@ def resolve_multiplex_source_config( raise ValueError( "source OME channel metadata changed after multiplex request resolution" ) + native_selection: list[dict[str, Any]] = [] + if native_multiplex: + from raw2features.multiplex.panel import resolve_marker_selection + + native_selection = resolve_marker_selection(names, cfg.multiplex_markers) + if ( + cfg.resolved_native_marker_selection + and native_selection != cfg.resolved_native_marker_selection + ): + raise ValueError( + "source marker selection changed after multiplex request resolution" + ) nuclear_indices: list[int] = [] if not cfg.no_seg: from raw2features.segmenters.nuclear import NuclearSegmenter @@ -740,6 +985,7 @@ def resolve_multiplex_source_config( return replace( cfg, resolved_channel_names=names, + resolved_native_marker_selection=native_selection, resolved_nuclear_channel_indices=nuclear_indices, resolved_original_channel_names=original_names, ) @@ -1093,6 +1339,7 @@ def run_slide( cfg, model_specs=injected_specs, ) + _validate_native_patch_geometry(cfg, injected_specs) import torch # noqa: F401 - imported so a torch-less env fails clearly, here @@ -1298,7 +1545,12 @@ def embedder_factory(device: str) -> list: ) run_contracts = {name: model_contracts[name] for name in models_to_do} _assert_loaded_model_contracts( - run_embedders, run_contracts, cfg.resolved_channel_names + run_embedders, + run_contracts, + cfg.resolved_channel_names, + selected_channels=cfg.resolved_native_marker_selection, + explicit_selection=bool(cfg.multiplex_markers), + native_model_params=cfg.native_multiplex_params, ) # Modality of this run: a marker stack routes the # nuclear segmenter + N-channel reads; brightfield is the RGB default. @@ -1343,7 +1595,13 @@ def embedder_factory(device: str) -> list: if multiplex: for e in run_embedders: panel_meta[e.name] = redact_metadata_credentials( - e.set_panel(reader.channel_names) + e.bind_multiplex_panel( + list(reader.channel_names or []), + selected_channels=cfg.resolved_native_marker_selection, + model_params=cfg.native_multiplex_params.get( + e.name, {} + ), + ) ) if append: ( @@ -1776,7 +2034,13 @@ def resolve_run( models=list(group.models), target_mpp=group.mpp, patch_px=group.patch_px, + native_multiplex_params={ + model: cfg.native_multiplex_params[model] + for model in group.models + if model in cfg.native_multiplex_params + }, ) + _validate_native_patch_geometry(group_cfg, model_specs) if not group_cfg.no_seg and _models_need_nuclear( list(group.models), model_specs ): @@ -1881,6 +2145,9 @@ def embed_slide( cfg.models, model_specs, cfg.resolved_channel_names, + selected_channels=cfg.resolved_native_marker_selection, + explicit_selection=bool(cfg.multiplex_markers), + native_model_params=cfg.native_multiplex_params, ) expected_grid_models = { group_cfg.grid_hash(): list(group_cfg.models) for group_cfg in group_cfgs @@ -2071,6 +2338,9 @@ def expected_model_contracts( cfg.models, specs, cfg.resolved_channel_names, + selected_channels=cfg.resolved_native_marker_selection, + explicit_selection=bool(cfg.multiplex_markers), + native_model_params=cfg.native_multiplex_params, ) @@ -2079,6 +2349,10 @@ def _bind_native_multiplex_panel_contracts( models: list[str], specs: dict, channel_names: list[str], + *, + selected_channels: list[dict[str, Any]] | None = None, + explicit_selection: bool = False, + native_model_params: dict[str, dict[str, Any]] | None = None, ) -> dict[str, dict]: """Bind a positional source panel into every native multiplex fingerprint. @@ -2090,6 +2364,13 @@ def _bind_native_multiplex_panel_contracts( if not channel_names: return dict(contracts) + selection = list(selected_channels or []) + if not selection: + selection = [ + {"source_index": index, "source_name": name} + for index, name in enumerate(channel_names) + ] + params_by_model = dict(native_model_params or {}) bound = {name: dict(contract) for name, contract in contracts.items()} for name in models: spec = specs.get(name) or get_spec(name) @@ -2100,12 +2381,33 @@ def _bind_native_multiplex_panel_contracts( continue fingerprint = bound[name]["output_fingerprint"] payload = json.loads(json.dumps(fingerprint["payload"])) - payload["multiplex_panel"] = { - "binding_contract_version": 1, - "channel_axis": "c", - "physical_channel_count": len(channel_names), - "effective_channel_names": list(channel_names), - } + model_params = params_by_model.get(name, {}) + use_v2 = ( + getattr(spec, "family", "") == "kronos2" + or explicit_selection + or bool(model_params) + ) + if use_v2: + payload["multiplex_panel"] = { + "binding_contract_version": 2, + "channel_axis": "c", + "physical_channel_count": len(channel_names), + "effective_channel_names": list(channel_names), + "selection": { + "mode": "explicit" if explicit_selection else "all", + "channels": json.loads(json.dumps(selection)), + }, + "model_params": json.loads(json.dumps(model_params)), + } + else: + # Preserve the released KRONOSv1 all-channel fingerprint exactly. Its + # historical contract already binds the complete positional source panel. + payload["multiplex_panel"] = { + "binding_contract_version": 1, + "channel_axis": "c", + "physical_channel_count": len(channel_names), + "effective_channel_names": list(channel_names), + } bound[name]["output_fingerprint"] = make_output_fingerprint(payload) return bound @@ -2182,7 +2484,14 @@ def _probe_factory_contracts( probe = factory(worker_devices[0]) try: contracts = _expected_contracts_for_devices(cfg, probe, worker_devices) - _assert_loaded_model_contracts(probe, contracts, cfg.resolved_channel_names) + _assert_loaded_model_contracts( + probe, + contracts, + cfg.resolved_channel_names, + selected_channels=cfg.resolved_native_marker_selection, + explicit_selection=bool(cfg.multiplex_markers), + native_model_params=cfg.native_multiplex_params, + ) specs = {embedder.name: embedder.spec for embedder in probe} return specs, contracts finally: @@ -2204,7 +2513,12 @@ def _amp_label(dtype) -> str: def _loaded_model_contracts( - embedders: list, channel_names: list[str] | None = None + embedders: list, + channel_names: list[str] | None = None, + *, + selected_channels: list[dict[str, Any]] | None = None, + explicit_selection: bool = False, + native_model_params: dict[str, dict[str, Any]] | None = None, ) -> dict[str, dict]: """Derive contracts from loaded specs and their effective execution precision.""" @@ -2225,6 +2539,9 @@ def _loaded_model_contracts( [embedder.name for embedder in embedders], {embedder.name: embedder.spec for embedder in embedders}, list(channel_names or []), + selected_channels=selected_channels, + explicit_selection=explicit_selection, + native_model_params=native_model_params, ) @@ -2232,10 +2549,20 @@ def _assert_loaded_model_contracts( embedders: list, expected: dict[str, dict], channel_names: list[str] | None = None, + *, + selected_channels: list[dict[str, Any]] | None = None, + explicit_selection: bool = False, + native_model_params: dict[str, dict[str, Any]] | None = None, ) -> None: """Fail before store mutation if loaded model copies differ from provenance.""" - actual = _loaded_model_contracts(embedders, channel_names) + actual = _loaded_model_contracts( + embedders, + channel_names, + selected_channels=selected_channels, + explicit_selection=explicit_selection, + native_model_params=native_model_params, + ) if set(actual) != set(expected): missing = sorted(set(expected) - set(actual)) extra = sorted(set(actual) - set(expected)) @@ -2579,6 +2906,22 @@ def _group_by_transform(embedders: list) -> list[list]: return list(groups.values()) +def _effective_batch_size(configured: int, embedders: list) -> int: + """Apply every model's optional runtime batch cap to one shared patch loop.""" + + limits: list[int] = [] + for embedder in embedders: + limit = embedder.max_batch_size + if limit is None: + continue + if isinstance(limit, bool) or not isinstance(limit, int) or limit <= 0: + raise ValueError( + f"{embedder.name}.max_batch_size must be a positive integer or None" + ) + limits.append(limit) + return min([configured, *limits]) + + def _embed_patches( reader, coords, @@ -2669,7 +3012,8 @@ def _run_batches( """ prof = prof or null_profiler() n = int(coords.shape[0]) - starts = list(range(0, n, cfg.batch_size)) + batch_size = _effective_batch_size(cfg.batch_size, embedders) + starts = list(range(0, n, batch_size)) # Group models by preprocessing signature once (it is static across batches): # the transform is computed once per group and shared across its members. @@ -2699,7 +3043,7 @@ def _producer() -> None: for start in starts: if stop.is_set(): break - batch = coords[start : start + cfg.batch_size] + batch = coords[start : start + batch_size] with prof.stage("read"): patches = _decode_batch( reader, @@ -2744,7 +3088,10 @@ def _producer() -> None: # Computed once per group (shared signature -> identical tensor) and # reused by every model in the group; embed_batch only reads it. with prof.stage("transform"): - batch_tensor = group[0].transform_batch(patches, device) + transform_patches = group[0].select_multiplex_channels(patches) + batch_tensor = group[0].transform_batch( + transform_patches, device + ) for emb in group: with prof.stage("gpu"): feats = ( @@ -2849,7 +3196,12 @@ def _worker(idx: int, device: str, lo: int, hi: int) -> None: embedders = embedder_factory(device) if expected_contracts is not None: _assert_loaded_model_contracts( - embedders, expected_contracts, cfg.resolved_channel_names + embedders, + expected_contracts, + cfg.resolved_channel_names, + selected_channels=cfg.resolved_native_marker_selection, + explicit_selection=bool(cfg.multiplex_markers), + native_model_params=cfg.native_multiplex_params, ) collector = _FeatureCollector(hi - lo, model_dims, cfg.features_dtype) # Own reader per worker: per-reader chunk cache, thread-safe zarr reads. diff --git a/tests/test_channelwise_runner.py b/tests/test_channelwise_runner.py index abb5bda..11a1503 100644 --- a/tests/test_channelwise_runner.py +++ b/tests/test_channelwise_runner.py @@ -643,4 +643,4 @@ def test_channelwise_rejects_learned_slide_encoder_before_embedding( ) def test_brightfield_rejects_orphaned_multiplex_options(kwargs): with pytest.raises(ValueError, match="require multiplex_strategy"): - RunConfig(models=["mock"], **kwargs) + RunConfig(models=["resnet50"], **kwargs) diff --git a/tests/test_cli_commands.py b/tests/test_cli_commands.py index 101099b..60a6496 100644 --- a/tests/test_cli_commands.py +++ b/tests/test_cli_commands.py @@ -220,6 +220,57 @@ def fake_embed_slide(slide, out_dir, cfg, **kwargs): _assert_multiplex_config(captured["cfg"]) +def test_embed_cli_threads_canonical_kronos2_additional_marker_contract( + tmp_path, monkeypatch +): + embed_module = import_module("raw2features.cli.embed") + captured = {} + + marker_csv = tmp_path / "novel.csv" + marker_csv.write_text( + "marker_name,marker_full_name,compartment,family,family_desc,mean,std\n" + "Novel-X,Novel marker X,membrane,immune_marker,immune marker,0.25,0.5\n" + ) + + def fake_embed_slide(slide, out_dir, cfg, **kwargs): + captured["cfg"] = cfg + return {"status": "complete"} + + monkeypatch.setattr(embed_module, "embed_slide", fake_embed_slide) + result = CliRunner().invoke( + app, + [ + "embed", + str(tmp_path / "unused.zarr"), + str(tmp_path / "out"), + "-m", + "kronos2", + "--no-seg", + "--marker", + "Novel-X", + "--kronos2-additional-markers", + str(marker_csv), + ], + ) + + assert result.exit_code == 0, result.output + contract = captured["cfg"].native_multiplex_params["kronos2"][ + "additional_markers" + ] + assert contract["rows"][0]["marker_name"] == "Novel-X" + assert "original_file_sha256" not in contract + + +@pytest.mark.parametrize("command", ["embed", "embed-many", "verify"]) +def test_kronos2_additional_marker_option_is_exposed_by_all_embedding_commands(command): + result = CliRunner().invoke( + app, [command, "--help"], env={"COLUMNS": "160"} + ) + + assert result.exit_code == 0, result.output + assert "--kronos2-additional-markers" in strip_ansi(result.output) + + def test_embed_many_cli_threads_multiplex_options(tmp_path, monkeypatch): embed_many_module = import_module("raw2features.cli.embed_many") slides = tmp_path / "slides" diff --git a/tests/test_kronos2.py b/tests/test_kronos2.py new file mode 100644 index 0000000..1c20e60 --- /dev/null +++ b/tests/test_kronos2.py @@ -0,0 +1,910 @@ +"""KRONOS2's pinned loader, marker identity, and output contract.""" + +from __future__ import annotations + +import hashlib +import json +import sys +from dataclasses import replace +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import numpy as np +import pytest + +from raw2features.embedders.fingerprint import ( + patch_output_fingerprint, + resolved_patch_amp, +) +from raw2features.embedders.kronos2_embedder import ( + KRONOS2_ADDITIONAL_MARKER_REGISTRATION_CONTRACT, + KRONOS2_BIOLINKBERT_ARTIFACT_SHA256, + KRONOS2_BIOLINKBERT_REVISION, + KRONOS2_BIOLINKBERT_SOURCE, + KRONOS2_FORWARD_CONTRACT, + KRONOS2_MARKER_MATCHING_CONTRACT, + KRONOS2_SCALING_CONTRACT, + KRONOS2_SNAPSHOT_ALLOW_PATTERNS, + KRONOS2_SNAPSHOT_TOP_LEVEL_ALLOWLIST, + Kronos2Embedder, + _marker_match_key, + _scaling_factor, +) +from raw2features.embedders.kronos2_metadata import ( + KRONOS2_ADDITIONAL_MARKER_COLUMNS, + parse_kronos2_additional_markers, +) +from raw2features.embedders.model_registry import get_spec +from raw2features.pipeline.runner import RunConfig, expected_model_contracts + + +def _digest(data: bytes) -> str: + return hashlib.sha256(data).hexdigest() + + +def _record( + canonical: str, + *, + pretraining: bool = True, + family: str = "immune_marker", +) -> dict: + return { + "canonical": canonical, + "pretraining": pretraining, + "compartment": "membrane", + "family": family, + } + + +def _additional_contract( + tmp_path, + *, + marker_name: str = "Novel-X", + marker_full_name: str = "Novel marker X", + compartment: str = "membrane", + family: str = "immune_marker", + family_desc: str = "immune marker", + mean: float = 0.25, + std: float = 0.5, + filename: str = "additional.csv", +): + path = tmp_path / filename + path.write_text( + ",".join(KRONOS2_ADDITIONAL_MARKER_COLUMNS) + + "\n" + + ",".join( + [ + marker_name, + marker_full_name, + compartment, + family, + family_desc, + str(mean), + str(std), + ] + ) + + "\n", + encoding="utf-8", + ) + return parse_kronos2_additional_markers(path) + + +def _panel_only_embedder() -> Kronos2Embedder: + emb = Kronos2Embedder(get_spec("kronos2")) + base = { + _marker_match_key("CD3"): _record("CD3"), + _marker_match_key("DAPI"): _record("DAPI", family="dna_stain") + | {"compartment": "nucleus"}, + } + emb._base_metadata_index = base + emb._metadata_index = dict(base) + emb._base_compartments = frozenset({"membrane", "nucleus"}) + emb._base_families = frozenset({"immune_marker", "dna_stain"}) + return emb + + +def test_kronos2_registry_and_fingerprint_cover_the_complete_static_contract(): + spec = get_spec("kronos2") + assert spec.family == "kronos2" + assert spec.modality == "multiplex" + assert spec.source == "hf_hub:MahmoodLab/KRONOS2" + assert spec.embedding_dim == 768 + assert spec.input_size == spec.extract_px == 256 + assert spec.pooling == "cls" + assert spec.inference_amp == "fp32" + assert spec.weights_revision == "e2a5473206c02d9973885b6445ef783f7e082951" + assert spec.weights_sha256 == ( + "afa261ef8b63bb0acea62612d7594ac11fbdb89fbdd23c513f046a034725d2b4" + ) + assert spec.timm_kwargs["marker_metadata_sha256"] == ( + "1e6791448a792c0edec9b2ba74efceaefba84d96852351cdc58d67f80b539074" + ) + + payload = patch_output_fingerprint(spec, "fp32")["payload"] + assert payload["checkpoint"]["effective"] == { + "repo": "MahmoodLab/KRONOS2", + "filename": "kronos2_vitb16_teacher.pth", + "mechanism": "pinned_local_snapshot", + } + constructor = payload["loader"]["constructor"] + assert constructor["input"] == "app_owned_allowlisted_pinned_local_snapshot" + assert constructor["transformers_local_files_only"] is True + assert constructor["upstream_nested_download_guard"] == "local_directory_input" + assert ( + constructor["upstream_sys_path_cleanup"] + == "remove_new_entry_after_construction" + ) + assert constructor["trust_remote_code"] is True + assert constructor["snapshot_allow_patterns"] == list( + KRONOS2_SNAPSHOT_ALLOW_PATTERNS + ) + assert constructor["snapshot_top_level_allowlist"] == list( + KRONOS2_SNAPSHOT_TOP_LEVEL_ALLOWLIST + ) + assert "*.py" not in KRONOS2_SNAPSHOT_ALLOW_PATTERNS + assert "test.py" not in KRONOS2_SNAPSHOT_TOP_LEVEL_ALLOWLIST + assert "sp_image.py" not in KRONOS2_SNAPSHOT_TOP_LEVEL_ALLOWLIST + assert constructor["marker_matching"] == KRONOS2_MARKER_MATCHING_CONTRACT + assert constructor["preferred_dapi_policy"] == ( + "canonical_dapi_else_draq5_else_none" + ) + assert constructor["novel_marker_registration"] == ( + "explicit_additional_markers_parameter_with_conditional_contract" + ) + assert constructor["forward"] == { + **KRONOS2_FORWARD_CONTRACT, + "method": "__call__", + "marker_names": "canonical_marker_metadata_names", + "result": "x_norm_clstoken", + } + assert constructor["marker_metadata"] == { + "filename": "marker_metadata.csv", + "sha256": spec.timm_kwargs["marker_metadata_sha256"], + } + assert constructor["published_marker_text_embeddings"] == "checkpoint_buffer" + assert payload["preprocessing"]["native_multiplex"] == KRONOS2_SCALING_CONTRACT + assert KRONOS2_FORWARD_CONTRACT["forward_microbatch_size"] == 8 + assert payload["output"] == { + "pooling": "cls", + "embedding_dim": 768, + "reg_tokens": 0, + "modality": "multiplex", + "resolved_amp": "fp32", + } + assert KRONOS2_BIOLINKBERT_SOURCE not in json.dumps(payload, sort_keys=True) + with pytest.raises(ValueError, match="does not accept requested AMP"): + resolved_patch_amp(spec, "fp16", "cuda") + with pytest.raises(ValueError, match="does not accept requested AMP"): + resolved_patch_amp(spec, "bf16", "cuda") + + +def test_novel_marker_contract_conditionally_enters_the_output_fingerprint(tmp_path): + additional = _additional_contract(tmp_path) + cfg = RunConfig( + models=["kronos2"], + no_seg=True, + device="cpu", + resolved_channel_names=["DAPI", "Novel-X"], + resolved_native_marker_selection=[ + {"source_index": 0, "source_name": "DAPI"}, + {"source_index": 1, "source_name": "Novel-X"}, + ], + native_multiplex_params={ + "kronos2": {"additional_markers": additional} + }, + ) + + payload = expected_model_contracts(cfg)["kronos2"]["output_fingerprint"][ + "payload" + ] + dynamic = payload["multiplex_panel"]["model_params"]["additional_markers"] + + assert dynamic["content_sha256"] == additional["content_sha256"] + assert dynamic["registration"]["text_encoder"]["source"] == ( + KRONOS2_BIOLINKBERT_SOURCE + ) + assert dynamic["registration"]["text_encoder"]["revision"] == ( + KRONOS2_BIOLINKBERT_REVISION + ) + assert dynamic["registration"]["text_encoder"]["artifact_sha256"] == ( + KRONOS2_BIOLINKBERT_ARTIFACT_SHA256 + ) + + +def test_marker_matching_is_separator_only_and_never_uses_v1_aliases(): + assert _marker_match_key("CD-8") == _marker_match_key("CD_8") == "cd8" + assert _marker_match_key("TCR-Vα7.2") == "tcrva72" + + emb = Kronos2Embedder(get_spec("kronos2")) + emb._metadata_index = { + _marker_match_key("CD8"): _record("CD8"), + _marker_match_key("panCK"): _record("panCK"), + _marker_match_key("CYTOKERATIN"): _record("CYTOKERATIN"), + _marker_match_key("DAPI"): _record("DAPI", family="dna_stain"), + } + summary = emb.set_panel(["CD-8", "panCK", "DAPI", ""]) + assert summary["matched"] == ["CD8", "panCK", "DAPI"] + assert summary["preferred_dapi"] == "DAPI" + assert summary["dropped"] == [""] + assert [row["channel_index"] for row in summary["mapping"]] == [0, 1, 2] + + # KRONOSv1 aliases CK -> CYTOKERATIN and Hoechst -> DAPI. KRONOS2 must not. + with pytest.raises(ValueError, match=r"\['CK', 'Hoechst1'\]"): + emb.set_panel(["CD8", "CK", "Hoechst1"]) + + +def test_panel_provenance_marks_published_non_pretraining_entries(): + spec = get_spec("kronos2") + emb = Kronos2Embedder(spec) + emb._metadata_index = { + _marker_match_key("CD3"): _record("CD3"), + _marker_match_key("DRAQ5"): _record( + "DRAQ5", pretraining=False, family="dna_stain" + ), + } + + summary = emb.set_panel(["CD3", "DRAQ-5"]) + assert summary["matched"] == ["CD3", "DRAQ5"] + assert summary["non_pretraining"] == ["DRAQ5"] + assert summary["preferred_dapi"] == "DRAQ5" + assert summary["defaulted"] == [] + assert summary["vocabulary"] == { + "source": spec.source, + "filename": "marker_metadata.csv", + "revision": spec.weights_revision, + "sha256": spec.timm_kwargs["marker_metadata_sha256"], + } + assert summary["novel_markers"] == { + "upstream_registration_available": True, + "enabled": False, + } + + +def test_normal_panel_binding_never_resolves_biolinkbert(monkeypatch): + import raw2features.embedders.kronos2_embedder as module + + emb = _panel_only_embedder() + monkeypatch.setattr( + module, + "download_pinned_hf_snapshot", + lambda *_args, **_kwargs: pytest.fail( + "ordinary published-marker binding downloaded BioLinkBERT" + ), + ) + + summary = emb.bind_multiplex_panel( + ["DAPI", "CD3"], + selected_channels=[{"source_index": 1, "source_name": "CD3"}], + model_params={}, + ) + + assert summary["matched"] == ["CD3"] + assert summary["novel_markers"]["enabled"] is False + assert emb._raw2features_multiplex_source_indices == (1,) + + +def test_novel_registration_is_pinned_exact_ordered_and_instance_idempotent( + monkeypatch, tmp_path +): + import raw2features.embedders.kronos2_embedder as module + + torch = ModuleType("torch") + + class Device: + def __init__(self, name): + self.type = str(name) + + def __eq__(self, other): + return isinstance(other, Device) and self.type == other.type + + class InferenceMode: + def __enter__(self): + return None + + def __exit__(self, *_args): + return False + + class FakeTensor: + def __init__(self, array): + self._array = np.asarray(array, dtype=np.float32) + + @property + def shape(self): + return self._array.shape + + def detach(self): + return self + + def cpu(self): + return self + + def float(self): + return self + + def numpy(self): + return self._array + + torch.device = Device + torch.inference_mode = InferenceMode + monkeypatch.setitem(sys.modules, "torch", torch) + + contract = _additional_contract(tmp_path) + snapshot = tmp_path / "biolinkbert" + snapshot.mkdir() + for filename in KRONOS2_BIOLINKBERT_ARTIFACT_SHA256: + (snapshot / filename).write_bytes(filename.encode()) + + events = [] + monkeypatch.setattr( + module, + "download_pinned_hf_snapshot", + lambda source, revision, *, allow_patterns: ( + events.append(("download", source, revision, allow_patterns)) + or str(snapshot) + ), + ) + monkeypatch.setattr( + module, + "verify_sha256", + lambda path, sha256, *, what: events.append( + ("verify", Path(path).name, sha256, what) + ), + ) + + builder_module_name = "_raw2features_test_kronos2_builder" + builder_module = ModuleType(builder_module_name) + encoder_inits = [] + prompt_inputs = [] + + class InnerEncoder: + def float(self): + events.append(("encoder_float",)) + return self + + def eval(self): + events.append(("encoder_eval",)) + return self + + def to(self, device): + events.append(("encoder_to", device.type)) + return self + + class BioLinkBERTTextEncoder: + def __init__(self, **kwargs): + events.append(("encoder_init",)) + encoder_inits.append(kwargs) + self.enc = InnerEncoder() + + def __call__(self, prompts): + events.append(("encode", tuple(prompts))) + return FakeTensor(np.ones((len(prompts), 1024), dtype=np.float32)) + + def build_prompt(*values): + prompt_inputs.append(values) + return "prompt::" + "::".join(values) + + builder_module.BioLinkBERTTextEncoder = BioLinkBERTTextEncoder + builder_module.build_prompt = build_prompt + monkeypatch.setitem(sys.modules, builder_module_name, builder_module) + Builder = type("MarkerEmbeddingBuilder", (), {}) + Builder.__module__ = builder_module_name + marker_builder = Builder() + original_cache = tmp_path / "original-cache" + marker_builder.cache_dir = original_cache + registration_calls = [] + temporary_dirs = [] + + def register_additional_markers(csv_path): + temporary_dirs.append(Path(marker_builder.cache_dir)) + array = np.load( + Path(marker_builder.cache_dir) / "additional_marker_text_embeddings.npy", + allow_pickle=False, + ) + registration_calls.append( + { + "csv": Path(csv_path).read_text(encoding="utf-8"), + "shape": tuple(array.shape), + "dtype": array.dtype, + } + ) + events.append(("register",)) + + emb = _panel_only_embedder() + emb._model = SimpleNamespace( + backbone=SimpleNamespace(marker_builder=marker_builder), + register_additional_markers=register_additional_markers, + ) + selection = [ + {"source_index": 2, "source_name": "Novel-X"}, + {"source_index": 0, "source_name": "DAPI"}, + ] + params = {"additional_markers": contract} + + summary = emb.bind_multiplex_panel( + ["DAPI", "CD3", "Novel-X"], + selected_channels=selection, + model_params=params, + ) + + assert events[0] == ( + "download", + KRONOS2_BIOLINKBERT_SOURCE, + KRONOS2_BIOLINKBERT_REVISION, + tuple(KRONOS2_BIOLINKBERT_ARTIFACT_SHA256), + ) + assert [event[1] for event in events if event[0] == "verify"] == list( + KRONOS2_BIOLINKBERT_ARTIFACT_SHA256 + ) + assert max( + index for index, event in enumerate(events) if event[0] == "verify" + ) < next(index for index, event in enumerate(events) if event[0] == "encoder_init") + assert encoder_inits == [ + { + "model_id": str(snapshot), + "device": torch.device("cpu"), + "cache_dir": encoder_inits[0]["cache_dir"], + "max_length": 128, + "fp16_on_cuda": False, + } + ] + assert Path(encoder_inits[0]["cache_dir"]).name == "transformers-cache" + assert prompt_inputs == [("novel-x", "Novel marker X", "membrane", "immune marker")] + assert registration_calls[0]["shape"] == (1, 1024) + assert registration_calls[0]["dtype"] == np.dtype("float32") + assert registration_calls[0]["csv"].splitlines()[0] == ",".join( + KRONOS2_ADDITIONAL_MARKER_COLUMNS + ) + assert marker_builder.cache_dir == original_cache + assert not temporary_dirs[0].exists() + assert summary["matched"] == ["Novel-X", "DAPI"] + assert summary["preferred_dapi"] == "DAPI" + assert summary["novel_markers"] == { + "upstream_registration_available": True, + "enabled": True, + "content_sha256": contract["content_sha256"], + "markers": ["Novel-X"], + "registration": KRONOS2_ADDITIONAL_MARKER_REGISTRATION_CONTRACT, + } + assert emb._raw2features_multiplex_source_indices == (2, 0) + full_patch = np.stack( + [ + np.full((16, 16), 10, dtype=np.uint16), + np.full((16, 16), 20, dtype=np.uint16), + np.full((16, 16), 30, dtype=np.uint16), + ], + axis=2, + ) + selected = emb.select_multiplex_channels([full_patch])[0] + assert selected[0, 0].tolist() == [30, 10] + + # The same semantic contract is a true no-op on a warm model. + second = emb.bind_multiplex_panel( + ["DAPI", "CD3", "Novel-X"], + selected_channels=selection, + model_params=params, + ) + assert second["matched"] == ["Novel-X", "DAPI"] + assert len(encoder_inits) == len(registration_calls) == 1 + + # Removing the parameter also removes the novel marker from the active + # raw2features vocabulary even though upstream buffers remain registered. + with pytest.raises(ValueError, match="Novel-X"): + emb.bind_multiplex_panel( + ["Novel-X"], + model_params={}, + ) + + changed = _additional_contract( + tmp_path, + mean=0.3, + filename="changed.csv", + ) + with pytest.raises(ValueError, match="cannot be reused with different contract"): + emb.bind_multiplex_panel( + ["Novel-X"], + model_params={"additional_markers": changed}, + ) + + +@pytest.mark.parametrize( + ("marker_name", "compartment", "family", "message"), + [ + ("CD-3", "membrane", "immune_marker", "collides"), + ("Novel-X", "cytoplasm", "immune_marker", "compartment"), + ("Novel-X", "membrane", "unknown_family", "family"), + ], +) +def test_novel_contract_collisions_and_categories_fail_before_biolink_execution( + monkeypatch, + tmp_path, + marker_name, + compartment, + family, + message, +): + import raw2features.embedders.kronos2_embedder as module + + contract = _additional_contract( + tmp_path, + marker_name=marker_name, + compartment=compartment, + family=family, + ) + emb = _panel_only_embedder() + monkeypatch.setattr( + module, + "download_pinned_hf_snapshot", + lambda *_args, **_kwargs: pytest.fail( + "BioLinkBERT resolved before marker metadata validation" + ), + ) + + with pytest.raises(ValueError, match=message): + emb.bind_multiplex_panel( + [marker_name], + model_params={"additional_markers": contract}, + ) + + +def test_only_published_dapi_or_draq5_gets_the_preferred_dapi_override(): + emb = _panel_only_embedder() + emb._metadata_index[_marker_match_key("NovelDNA")] = _record( + "NovelDNA", pretraining=False, family="dna_stain" + ) + emb._binding_native_panel = True + try: + summary = emb.set_panel(["NovelDNA"]) + finally: + emb._binding_native_panel = False + + assert summary["preferred_dapi"] is None + + +@pytest.mark.parametrize( + ("dtype", "expected"), + [ + (np.dtype("uint8"), 255.0), + (np.dtype("uint16"), 65535.0), + (np.dtype("uint32"), 65535.0), + (np.dtype("float32"), 400.0), + (np.dtype("float64"), 400.0), + ], +) +def test_kronos2_uses_the_upstream_dtype_divisor(dtype, expected): + assert _scaling_factor(dtype) == expected + + +def test_kronos2_rejects_a_dtype_without_published_scaling(): + with pytest.raises(ValueError, match="no published input scaling"): + _scaling_factor(np.dtype("int16")) + + +def test_transform_delegates_zscore_to_upstream_with_canonical_markers(): + torch = pytest.importorskip("torch") + emb = Kronos2Embedder(get_spec("kronos2")) + emb._metadata_index = { + _marker_match_key("CD3"): _record("CD3"), + _marker_match_key("DRAQ5"): _record( + "DRAQ5", pretraining=False, family="dna_stain" + ), + } + emb.set_panel(["CD-3", "DRAQ-5"]) + calls = [] + + class Model: + def preprocess(self, patches, names, *, preferred_dapi): + calls.append((patches.copy(), list(names), preferred_dapi)) + return patches + np.float32(1.0) + + emb._model = Model() + patches = [ + np.stack( + [ + np.full((16, 16), 32767, dtype=np.uint16), + np.full((16, 16), 65535, dtype=np.uint16), + ], + axis=2, + ) + ] + output = emb.transform_batch(patches, "cpu") + + scaled, names, preferred_dapi = calls[0] + assert scaled.dtype == np.float32 + assert tuple(scaled.shape) == (1, 2, 16, 16) + assert np.allclose(scaled[0, 0], np.float32(32767 / 65535)) + assert np.allclose(scaled[0, 1], 1.0) + assert names == ["CD3", "DRAQ5"] + assert preferred_dapi == "DRAQ5" + assert output.dtype == torch.float32 + assert output.device.type == "cpu" + assert np.allclose(output.numpy(), scaled + 1.0) + + +@pytest.mark.parametrize( + ("patch", "message"), + [ + (np.zeros((16, 16), dtype=np.uint16), "must be HWC"), + (np.zeros((16, 32, 1), dtype=np.uint16), "must be square"), + (np.zeros((15, 15, 1), dtype=np.uint16), "divisible by"), + (np.zeros((16, 16, 2), dtype=np.uint16), "bound selected panel"), + ( + np.full((16, 16, 1), np.nan, dtype=np.float32), + "non-finite float", + ), + ], +) +def test_transform_rejects_invalid_native_patch_contract(patch, message): + emb = Kronos2Embedder(get_spec("kronos2")) + emb._metadata_index = {_marker_match_key("CD3"): _record("CD3")} + emb.set_panel(["CD3"]) + emb._model = type( + "Model", + (), + {"preprocess": staticmethod(lambda patches, names, *, preferred_dapi: patches)}, + )() + + with pytest.raises(ValueError, match=message): + emb.transform_batch([patch], "cuda") + + +@pytest.mark.parametrize("n_rows", [1, 7, 8, 9]) +def test_forward_uses_padded_fixed_eight_row_microbatches(n_rows): + torch = pytest.importorskip("torch") + emb = Kronos2Embedder(get_spec("kronos2")) + emb._metadata_index = {_marker_match_key("CD3"): _record("CD3")} + emb.set_panel(["CD3"]) + emb._device = "cpu" + calls = [] + + class Model: + def __call__(self, batch, marker_names): + calls.append( + { + "shape": tuple(batch.shape), + "rows": batch[:, 0, 0, 0].detach().cpu().tolist(), + "marker_names": list(marker_names), + } + ) + values = batch[:, 0, 0, 0].reshape(8, 1) + return values.repeat(1, emb.embedding_dim) + + emb._model = Model() + batch = ( + torch.arange(n_rows, dtype=torch.float32) + .reshape(n_rows, 1, 1, 1) + .expand(n_rows, 1, 16, 16) + .clone() + ) + + output = emb.embed_batch(batch) + + assert emb.max_batch_size == 8 + assert tuple(output.shape) == (n_rows, emb.embedding_dim) + assert output[:, 0].tolist() == list(range(n_rows)) + assert all(call["shape"] == (8, 1, 16, 16) for call in calls) + assert all(call["marker_names"] == ["CD3"] for call in calls) + assert len(calls) == (n_rows + 7) // 8 + assert all(len(call["rows"]) == 8 for call in calls) + if n_rows % 8: + final_rows = calls[-1]["rows"] + valid = n_rows % 8 + assert final_rows[valid:] == [final_rows[valid - 1]] * (8 - valid) + + +def test_loader_verifies_both_artifacts_then_loads_only_the_local_snapshot( + monkeypatch, tmp_path +): + import raw2features.embedders.kronos2_embedder as module + + weight_bytes = b"pinned KRONOS2 weights" + metadata_bytes = ( + b",marker_name,compartment,family,mean,std,pretraining\n" + b"0,DAPI,nucleus,dna_stain,0.1,0.2,yes\n" + b"1,DRAQ5,nucleus,dna_stain,0.1,0.2,no\n" + ) + weights = tmp_path / "kronos2_vitb16_teacher.pth" + metadata = tmp_path / "marker_metadata.csv" + weights.write_bytes(weight_bytes) + metadata.write_bytes(metadata_bytes) + snapshot_calls = [] + load_calls = [] + events = [] + + isolated_dir = tmp_path / "isolated-runtime" + monkeypatch.setattr( + module, + "pinned_model_cache_dir", + lambda source, revision: str(isolated_dir), + ) + + def snapshot_download(source, revision, *, allow_patterns, local_dir): + snapshot_calls.append((source, revision, allow_patterns, local_dir)) + return str(tmp_path) + + monkeypatch.setattr(module, "download_pinned_hf_snapshot", snapshot_download) + original_verify = module.verify_sha256 + + def verify(path, expected, *, what): + events.append(("verify", Path(path).name, what)) + original_verify(path, expected, what=what) + + monkeypatch.setattr(module, "verify_sha256", verify) + + class Model: + def float(self): + return self + + def eval(self): + return self + + def to(self, device): + self.device = device + return self + + transformers = ModuleType("transformers") + + class AutoModel: + @staticmethod + def from_pretrained(source, **kwargs): + events.append(("load", Path(source).name, "AutoModel")) + load_calls.append((source, kwargs)) + sys.path.insert(0, source) + return Model() + + transformers.AutoModel = AutoModel + monkeypatch.setitem(sys.modules, "transformers", transformers) + torch = ModuleType("torch") + torch.float32 = object() + monkeypatch.setitem(sys.modules, "torch", torch) + + base = get_spec("kronos2") + spec = replace( + base, + weights_sha256=_digest(weight_bytes), + timm_kwargs={ + **base.timm_kwargs, + "marker_metadata_sha256": _digest(metadata_bytes), + }, + ) + emb = Kronos2Embedder(spec) + emb._registered_additional_markers_digest = "stale" + emb._novel_registration_failed = True + emb._raw2features_multiplex_source_indices = (9,) + emb.load(device="cpu") + + assert snapshot_calls == [ + ( + spec.source, + spec.weights_revision, + KRONOS2_SNAPSHOT_ALLOW_PATTERNS, + str(isolated_dir), + ) + ] + assert events == [ + ("verify", "kronos2_vitb16_teacher.pth", "kronos2"), + ("verify", "marker_metadata.csv", "kronos2 marker metadata"), + ("load", tmp_path.name, "AutoModel"), + ] + assert load_calls == [ + ( + str(tmp_path), + { + "trust_remote_code": True, + "local_files_only": True, + "device": "cpu", + }, + ) + ] + assert emb._dtype is torch.float32 + assert set(emb._metadata_index) == {"dapi", "draq5"} + assert not hasattr(emb, "_registered_additional_markers_digest") + assert not hasattr(emb, "_novel_registration_failed") + assert not hasattr(emb, "_raw2features_multiplex_source_indices") + assert str(tmp_path) not in sys.path + + +def test_loader_rejects_stale_unrequested_snapshot_code_before_execution( + monkeypatch, tmp_path +): + import raw2features.embedders.kronos2_embedder as module + + weight_bytes = b"pinned KRONOS2 weights" + metadata_bytes = ( + b",marker_name,compartment,family,mean,std,pretraining\n" + b"0,DAPI,nucleus,dna_stain,0.1,0.2,yes\n" + ) + (tmp_path / "kronos2_vitb16_teacher.pth").write_bytes(weight_bytes) + (tmp_path / "marker_metadata.csv").write_bytes(metadata_bytes) + (tmp_path / "test.py").write_text( + "raise AssertionError('must never execute')\n", encoding="utf-8" + ) + monkeypatch.setattr( + module, + "pinned_model_cache_dir", + lambda *_args: str(tmp_path / "isolated-runtime"), + ) + monkeypatch.setattr( + module, + "download_pinned_hf_snapshot", + lambda *_args, **_kwargs: str(tmp_path), + ) + + transformers = ModuleType("transformers") + + class AutoModel: + @staticmethod + def from_pretrained(*_args, **_kwargs): + pytest.fail("unexpected snapshot code reached AutoModel") + + transformers.AutoModel = AutoModel + monkeypatch.setitem(sys.modules, "transformers", transformers) + torch = ModuleType("torch") + torch.float32 = object() + monkeypatch.setitem(sys.modules, "torch", torch) + + base = get_spec("kronos2") + spec = replace( + base, + weights_sha256=_digest(weight_bytes), + timm_kwargs={ + **base.timm_kwargs, + "marker_metadata_sha256": _digest(metadata_bytes), + }, + ) + with pytest.raises(ValueError, match=r"unexpected.*test\.py"): + Kronos2Embedder(spec).load(device="cpu") + + +def test_loader_rejects_compile_before_import_or_download(monkeypatch): + import raw2features.embedders.kronos2_embedder as module + + monkeypatch.setattr( + module, + "download_pinned_hf_snapshot", + lambda *_args, **_kwargs: pytest.fail( + "KRONOS2 downloaded a snapshot before rejecting compile" + ), + ) + + with pytest.raises(ValueError, match="does not support --compile"): + Kronos2Embedder(get_spec("kronos2")).load(device="cpu", compile=True) + + +def test_loader_refuses_bad_marker_metadata_before_custom_code_executes( + monkeypatch, tmp_path +): + import raw2features.embedders.kronos2_embedder as module + + weights = tmp_path / "kronos2_vitb16_teacher.pth" + metadata = tmp_path / "marker_metadata.csv" + weights.write_bytes(b"weights") + metadata.write_text( + ",marker_name,compartment,family,mean,std,pretraining\n" + "0,DAPI,nucleus,dna_stain,0.1,0.2,yes\n" + ) + monkeypatch.setattr( + module, + "download_pinned_hf_snapshot", + lambda *_args, **_kwargs: str(tmp_path), + ) + torch = ModuleType("torch") + torch.float32 = object() + monkeypatch.setitem(sys.modules, "torch", torch) + transformers = ModuleType("transformers") + + class AutoModel: + @staticmethod + def from_pretrained(*_args, **_kwargs): + pytest.fail("custom model code ran before marker metadata verification") + + transformers.AutoModel = AutoModel + monkeypatch.setitem(sys.modules, "transformers", transformers) + base = get_spec("kronos2") + spec = replace( + base, + weights_sha256=_digest(b"weights"), + timm_kwargs={**base.timm_kwargs, "marker_metadata_sha256": "0" * 64}, + ) + + with pytest.raises(ValueError, match="marker metadata"): + Kronos2Embedder(spec).load(device="cpu") diff --git a/tests/test_kronos2_metadata.py b/tests/test_kronos2_metadata.py new file mode 100644 index 0000000..558a854 --- /dev/null +++ b/tests/test_kronos2_metadata.py @@ -0,0 +1,274 @@ +"""Canonical validation for user-supplied KRONOS2 marker metadata.""" + +from __future__ import annotations + +import hashlib +import json + +import pytest + +from raw2features.embedders.kronos2_metadata import ( + KRONOS2_ADDITIONAL_MARKER_COLUMNS, + KRONOS2_NOVEL_MARKER_REGISTRATION_CONTRACT, + Kronos2MarkerMetadataError, + materialize_kronos2_additional_markers_csv, + parse_kronos2_additional_markers, +) + +HEADER = ",".join(KRONOS2_ADDITIONAL_MARKER_COLUMNS) + + +def _write(tmp_path, text: str, name: str = "markers.csv"): + path = tmp_path / name + path.write_bytes(text.encode("utf-8")) + return path + + +def _row( + marker_name: str = "FOXP3", + marker_full_name: str = "Forkhead box P3", + compartment: str = "nucleus", + family: str = "transcription_factor", + family_desc: str = "transcription factor", + mean: str = "0.25", + std: str = "0.5", +) -> str: + return ",".join( + [ + marker_name, + marker_full_name, + compartment, + family, + family_desc, + mean, + std, + ] + ) + + +def test_parser_returns_canonical_json_safe_ordered_contract(tmp_path): + source = ( + "\ufeff" + + HEADER + + ",ignored\n" + + _row(" FOXP3 ", mean="-0", std="5e-1") + + ",one\n" + + _row("Novel-2", " Novel marker 2 ", mean="1.25", std="2") + + ",two\n" + ) + path = _write(tmp_path, source) + + contract = parse_kronos2_additional_markers(path) + + assert contract["contract_version"] == 1 + assert contract["consumed_columns"] == list(KRONOS2_ADDITIONAL_MARKER_COLUMNS) + assert [row["marker_name"] for row in contract["rows"]] == ["FOXP3", "Novel-2"] + assert contract["rows"][0]["mean"] == 0.0 + assert contract["rows"][0]["std"] == 0.5 + assert contract["rows"][1]["mean"] == 1.25 + assert contract["rows"][1]["std"] == 2.0 + assert "original_file_sha256" not in contract + assert contract["registration"] == KRONOS2_NOVEL_MARKER_REGISTRATION_CONTRACT + json.dumps(contract, allow_nan=False) + + unhashed = { + "contract_version": contract["contract_version"], + "consumed_columns": contract["consumed_columns"], + "registration": contract["registration"], + "rows": contract["rows"], + } + canonical = json.dumps( + unhashed, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + allow_nan=False, + ) + assert ( + contract["content_sha256"] + == hashlib.sha256(canonical.encode("utf-8")).hexdigest() + ) + + +def test_semantic_digest_ignores_csv_format_column_order_and_extra_columns(tmp_path): + first = _write( + tmp_path, + HEADER + "\r\n" + _row(mean="0.2500", std="0.500") + "\r\n", + "first.csv", + ) + second = _write( + tmp_path, + "std,unused,marker_name,family_desc,mean,family,compartment," + "marker_full_name\n" + "5e-1,ignored,FOXP3,transcription factor,2.5e-1," + "transcription_factor,nucleus,Forkhead box P3\n", + "second.csv", + ) + + one = parse_kronos2_additional_markers(first) + two = parse_kronos2_additional_markers(second) + + assert one["rows"] == two["rows"] + assert one["content_sha256"] == two["content_sha256"] + assert "original_file_sha256" not in one + assert "original_file_sha256" not in two + + +def test_prompt_unicode_is_preserved_as_semantic_content(tmp_path): + fullwidth = _write( + tmp_path, + f"{HEADER}\n{_row('ＦＯＸＰ３', 'Ｆｏｒｋｈｅａｄ box P3')}\n", + "fullwidth.csv", + ) + ascii_text = _write( + tmp_path, + f"{HEADER}\n{_row('FOXP3', 'Forkhead box P3')}\n", + "ascii.csv", + ) + + one = parse_kronos2_additional_markers(fullwidth) + two = parse_kronos2_additional_markers(ascii_text) + + assert one["rows"][0]["marker_name"] == "ＦＯＸＰ３" + assert one["rows"][0]["marker_full_name"] == "Ｆｏｒｋｈｅａｄ box P3" + assert one["content_sha256"] != two["content_sha256"] + + +def test_original_file_digest_is_available_only_as_explicit_non_identity_audit_data( + tmp_path, +): + path = _write(tmp_path, f"{HEADER}\n{_row()}\n") + + contract = parse_kronos2_additional_markers(path, include_original_file_sha256=True) + + assert ( + contract["original_file_sha256"] + == hashlib.sha256(path.read_bytes()).hexdigest() + ) + + +def test_row_order_is_part_of_the_semantic_digest(tmp_path): + one = _row("A", "Marker A") + two = _row("B", "Marker B") + forward = _write(tmp_path, f"{HEADER}\n{one}\n{two}\n", "forward.csv") + reverse = _write(tmp_path, f"{HEADER}\n{two}\n{one}\n", "reverse.csv") + + assert ( + parse_kronos2_additional_markers(forward)["content_sha256"] + != parse_kronos2_additional_markers(reverse)["content_sha256"] + ) + + +@pytest.mark.parametrize("missing", KRONOS2_ADDITIONAL_MARKER_COLUMNS) +def test_all_consumed_columns_are_required(tmp_path, missing): + columns = [ + column for column in KRONOS2_ADDITIONAL_MARKER_COLUMNS if column != missing + ] + path = _write(tmp_path, ",".join(columns) + "\n" + ",".join("x" for _ in columns)) + + with pytest.raises(Kronos2MarkerMetadataError, match="missing required columns"): + parse_kronos2_additional_markers(path) + + +def test_duplicate_headers_and_empty_tables_are_rejected(tmp_path): + duplicate = _write( + tmp_path, + HEADER + ",marker_name\n" + _row() + ",duplicate\n", + "duplicate.csv", + ) + empty = _write(tmp_path, HEADER + "\n", "empty.csv") + + with pytest.raises(Kronos2MarkerMetadataError, match="duplicate columns"): + parse_kronos2_additional_markers(duplicate) + with pytest.raises(Kronos2MarkerMetadataError, match="at least one marker row"): + parse_kronos2_additional_markers(empty) + + +@pytest.mark.parametrize( + "column", + ["marker_name", "marker_full_name", "compartment", "family", "family_desc"], +) +def test_identity_fields_must_not_be_blank(tmp_path, column): + values = dict( + zip(KRONOS2_ADDITIONAL_MARKER_COLUMNS, _row().split(","), strict=True) + ) + values[column] = " " + path = _write( + tmp_path, + HEADER + + "\n" + + ",".join(values[name] for name in KRONOS2_ADDITIONAL_MARKER_COLUMNS), + ) + + with pytest.raises( + Kronos2MarkerMetadataError, match=rf"{column!r} must not be blank" + ): + parse_kronos2_additional_markers(path) + + +def test_control_characters_are_rejected_before_whitespace_is_trimmed(tmp_path): + path = _write( + tmp_path, + f"{HEADER}\n" + '"FOXP3","Forkhead\nbox P3",nucleus,transcription_factor,' + "transcription factor,0.25,0.5\n", + ) + + with pytest.raises(Kronos2MarkerMetadataError, match="control character"): + parse_kronos2_additional_markers(path) + + +@pytest.mark.parametrize( + ("mean", "std", "message"), + [ + ("nan", "1", "finite number"), + ("inf", "1", "finite number"), + ("not-a-number", "1", "finite number"), + ("1", "nan", "finite number"), + ("1", "0", "greater than zero"), + ("1", "-0.01", "greater than zero"), + ], +) +def test_statistics_must_be_finite_and_std_positive(tmp_path, mean, std, message): + path = _write(tmp_path, f"{HEADER}\n{_row(mean=mean, std=std)}\n") + + with pytest.raises(Kronos2MarkerMetadataError, match=message): + parse_kronos2_additional_markers(path) + + +@pytest.mark.parametrize( + ("first", "second", "message"), + [ + ("FOXP3", "foxp3", "case-insensitively"), + ("CD-8", "CD_8", "separator-insensitive"), + ("TCR-Vα7.2", "TCR_VA7.2", "separator-insensitive"), + ], +) +def test_duplicate_marker_identities_are_rejected(tmp_path, first, second, message): + path = _write( + tmp_path, + f"{HEADER}\n{_row(first, 'First')}\n{_row(second, 'Second')}\n", + ) + + with pytest.raises(Kronos2MarkerMetadataError, match=message): + parse_kronos2_additional_markers(path) + + +def test_materialized_csv_contains_only_consumed_columns_and_round_trips(tmp_path): + source = _write( + tmp_path, + HEADER + ",ignored\n" + _row("Novel A", "Novel A full") + ",value\n", + "source.csv", + ) + contract = parse_kronos2_additional_markers( + source, include_original_file_sha256=False + ) + destination = tmp_path / "canonical.csv" + + materialize_kronos2_additional_markers_csv(contract, destination) + + assert destination.read_text().splitlines()[0] == HEADER + reparsed = parse_kronos2_additional_markers(destination) + assert reparsed["rows"] == contract["rows"] + assert reparsed["content_sha256"] == contract["content_sha256"] + assert "original_file_sha256" not in contract diff --git a/tests/test_native_multiplex_selection.py b/tests/test_native_multiplex_selection.py new file mode 100644 index 0000000..e3938fc --- /dev/null +++ b/tests/test_native_multiplex_selection.py @@ -0,0 +1,366 @@ +"""Native multiplex marker selection, identity, and execution-safety contracts.""" + +from __future__ import annotations + +from dataclasses import replace + +import numpy as np +import pytest + +from raw2features.embedders.base import Embedder, ModelSpec +from raw2features.multiplex.panel import resolve_marker_selection +from raw2features.pipeline.runner import ( + RunConfig, + _effective_batch_size, + _reject_mixed_native_and_brightfield_models, + expected_model_contracts, + resolve_multiplex_source_config, + resolve_run, + run_slide, +) + + +def _native_spec(name: str) -> ModelSpec: + return ModelSpec( + name=name, + family="test", + source=f"hf_hub:example/{name}", + embedding_dim=4, + input_size=32, + pooling="cls", + mean=(0.0, 0.0, 0.0), + std=(1.0, 1.0, 1.0), + transform_source_url="https://example.org/model", + license="MIT", + gated=False, + modality="multiplex", + ) + + +class _NativeEmbedder(Embedder): + def __init__(self, name: str): + super().__init__(_native_spec(name)) + self.bound_names = None + + def set_panel(self, channel_names): + self.bound_names = list(channel_names or []) + return {"kept": self.bound_names} + + def load(self, device="cpu", dtype=None, compile=False): + return self + + def embed_batch(self, batch): # pragma: no cover - selection is tested pre-forward + raise AssertionError("not used") + + +class _LimitedNativeEmbedder(_NativeEmbedder): + def __init__(self, name: str, limit): + super().__init__(name) + self.limit = limit + + @property + def max_batch_size(self): + return self.limit + + +def test_selection_is_normalized_but_keeps_requested_order_and_physical_indices(): + panel = ["DAPI", "CD3", "PanCK"] + assert resolve_marker_selection(panel, [" pAnCk ", "dapi"]) == [ + {"source_index": 2, "source_name": "PanCK"}, + {"source_index": 0, "source_name": "DAPI"}, + ] + assert resolve_marker_selection(["DAPI", "", "CD3"]) == [ + {"source_index": 0, "source_name": "DAPI"}, + {"source_index": 1, "source_name": ""}, + {"source_index": 2, "source_name": "CD3"}, + ] + + +@pytest.mark.parametrize( + ("panel", "requested", "message"), + [ + (["CD3"], ["CD8"], "not present"), + (["CD3", "cd3"], ["CD3"], "ambiguous"), + (["CD3"], ["CD3", "cd3"], "more than once"), + ], +) +def test_selection_refuses_missing_ambiguous_or_duplicate_requests( + panel, requested, message +): + with pytest.raises(ValueError, match=message): + resolve_marker_selection(panel, requested) + + +def test_native_binding_slices_full_hwc_patches_by_physical_index_in_request_order(): + embedder = _NativeEmbedder("native") + selection = [ + {"source_index": 2, "source_name": "C"}, + {"source_index": 0, "source_name": "A"}, + ] + summary = embedder.bind_multiplex_panel( + ["A", "B", "C"], selected_channels=selection + ) + patch = np.stack( + [ + np.full((2, 2), 1, dtype=np.uint16), + np.full((2, 2), 2, dtype=np.uint16), + np.full((2, 2), 3, dtype=np.uint16), + ], + axis=2, + ) + selected = embedder.select_multiplex_channels([patch])[0] + + assert summary == { + "kept": ["C", "A"], + "source_selection": selection, + } + assert embedder.bound_names == ["C", "A"] + assert selected.shape == (2, 2, 2) + assert np.all(selected[..., 0] == 3) + assert np.all(selected[..., 1] == 1) + + +def test_kronos_v1_explicit_selection_provenance_uses_physical_source_indices(): + from raw2features.embedders.kronos_embedder import KronosEmbedder + from raw2features.embedders.model_registry import get_spec + + embedder = KronosEmbedder(get_spec("kronos")) + embedder._vocab = { + "DAPI": (4, 0.1, 0.2, "DAPI"), + "CD3": (3, 0.1, 0.2, "CD3"), + } + selection = [ + {"source_index": 2, "source_name": "CD3"}, + {"source_index": 0, "source_name": "DAPI"}, + ] + summary = embedder.bind_multiplex_panel( + ["DAPI", "ignored", "CD3"], selected_channels=selection + ) + + assert [row["channel_index"] for row in summary["mapping"]] == [2, 0] + assert [row["channel"] for row in summary["mapping"]] == ["CD3", "DAPI"] + assert summary["source_selection"] == selection + assert embedder._panel["idx"].tolist() == [0, 1] + + +def test_native_transform_signatures_never_cross_model_names(): + first = _NativeEmbedder("native_a") + second = _NativeEmbedder("native_b") + assert first.transform_signature != second.transform_signature + + +def test_model_batch_caps_lower_the_runtime_batch_and_fail_on_invalid_limits(): + ordinary = _NativeEmbedder("ordinary") + limited = _LimitedNativeEmbedder("limited", 8) + assert _effective_batch_size(256, [ordinary]) == 256 + assert _effective_batch_size(256, [ordinary, limited]) == 8 + assert _effective_batch_size(4, [limited]) == 4 + with pytest.raises(ValueError, match="positive integer or None"): + _effective_batch_size(256, [_LimitedNativeEmbedder("bad", 0)]) + + +def test_native_marker_selection_changes_output_identity_not_grid_identity(): + base = RunConfig( + models=["kronos"], + no_seg=True, + resolved_channel_names=["DAPI", "CD3", "CD8"], + resolved_native_marker_selection=[ + {"source_index": 0, "source_name": "DAPI"}, + {"source_index": 1, "source_name": "CD3"}, + {"source_index": 2, "source_name": "CD8"}, + ], + device="cpu", + ) + selected = replace( + base, + multiplex_markers=["CD8", "DAPI"], + resolved_native_marker_selection=[ + {"source_index": 2, "source_name": "CD8"}, + {"source_index": 0, "source_name": "DAPI"}, + ], + ) + + old_payload = expected_model_contracts(base)["kronos"]["output_fingerprint"][ + "payload" + ]["multiplex_panel"] + new_payload = expected_model_contracts(selected)["kronos"]["output_fingerprint"][ + "payload" + ]["multiplex_panel"] + assert old_payload == { + "binding_contract_version": 1, + "channel_axis": "c", + "physical_channel_count": 3, + "effective_channel_names": ["DAPI", "CD3", "CD8"], + } + assert new_payload["binding_contract_version"] == 2 + assert new_payload["selection"] == { + "mode": "explicit", + "channels": [ + {"source_index": 2, "source_name": "CD8"}, + {"source_index": 0, "source_name": "DAPI"}, + ], + } + assert base.grid_hash() == selected.grid_hash() + assert base.content_hash() != selected.content_hash() + + +def test_kronos2_uses_v2_panel_contract_even_for_the_all_channel_default(): + cfg = RunConfig( + models=["kronos2"], + no_seg=True, + resolved_channel_names=["DAPI", "CD3"], + resolved_native_marker_selection=[ + {"source_index": 0, "source_name": "DAPI"}, + {"source_index": 1, "source_name": "CD3"}, + ], + device="cpu", + ) + panel = expected_model_contracts(cfg)["kronos2"]["output_fingerprint"]["payload"][ + "multiplex_panel" + ] + assert panel["binding_contract_version"] == 2 + assert panel["selection"]["mode"] == "all" + assert panel["selection"]["channels"] == cfg.resolved_native_marker_selection + assert panel["model_params"] == {} + + +@pytest.mark.parametrize("amp", ["fp16", "bf16"]) +def test_kronos2_rejects_explicit_lower_precision_before_source_work(amp): + with pytest.raises(ValueError, match="published fp32 inference path"): + RunConfig(models=["kronos2"], amp=amp) + + +def test_kronos2_rejects_compile_before_source_work(): + with pytest.raises(ValueError, match="does not support --compile"): + RunConfig(models=["kronos2"], compile=True) + + +def test_kronos2_legacy_grid_discovery_does_not_build_forbidden_amp_requests(): + no_seg = RunConfig(models=["kronos2"], no_seg=True, patch_px=256) + assert no_seg.compatible_legacy_grid_hashes() + + segmented = RunConfig( + models=["kronos2"], + patch_px=256, + resolved_channel_names=["DAPI", "CD3"], + resolved_nuclear_channel_indices=[0], + resolved_original_channel_names=["DAPI", "CD3"], + ) + hashes = segmented.compatible_legacy_grid_hashes() + evidence = segmented.compatible_legacy_grid_segmenters() + assert hashes + assert evidence + assert set(evidence).issubset(hashes) + assert set(evidence.values()) == {"nuclear"} + + +@pytest.mark.parametrize("patch_px", [64, 256]) +def test_kronos2_accepts_stride_aligned_patch_geometry(patch_px): + _groups, group_cfgs, _run_hash = resolve_run( + RunConfig(models=["kronos2"], no_seg=True, patch_px=patch_px), + requested_patch_px=patch_px, + ) + assert group_cfgs[0].patch_px == patch_px + + +def test_kronos2_rejects_unaligned_patch_geometry_before_execution(): + with pytest.raises(ValueError, match="divisible by 16"): + resolve_run( + RunConfig(models=["kronos2"], no_seg=True, patch_px=250), + requested_patch_px=250, + ) + + +def test_single_grid_primitive_rejects_unaligned_kronos2_before_store_mutation( + synthetic_multiplex_ngff, tmp_path +): + out_dir = tmp_path / "out" + + with pytest.raises(ValueError, match="divisible by 16"): + run_slide( + synthetic_multiplex_ngff, + str(out_dir), + RunConfig( + models=["kronos2"], + no_seg=True, + patch_px=250, + device="cpu", + ), + ) + + assert not out_dir.exists() + + +def test_source_resolution_selects_native_markers_but_keeps_nuclear_full_panel( + synthetic_multiplex_ngff, +): + cfg = RunConfig( + models=["kronos"], + multiplex_markers=["foxp3", "DAPI"], + device="cpu", + ) + resolved = resolve_multiplex_source_config(synthetic_multiplex_ngff, cfg) + assert resolved.resolved_native_marker_selection == [ + {"source_index": 4, "source_name": "FOXP3"}, + {"source_index": 0, "source_name": "DAPI"}, + ] + assert resolved.resolved_channel_names == ["DAPI", "CD3", "CD8", "CD20", "FOXP3"] + assert resolved.resolved_nuclear_channel_indices == [0] + + +def test_injected_native_model_can_use_marker_selection( + synthetic_multiplex_ngff, +): + spec = _native_spec("external_native") + cfg = RunConfig( + models=["external_native"], + multiplex_markers=["FOXP3", "DAPI"], + device="cpu", + ) + + resolved = resolve_multiplex_source_config( + synthetic_multiplex_ngff, + cfg, + model_specs={"external_native": spec}, + ) + + assert resolved.resolved_native_marker_selection == [ + {"source_index": 4, "source_name": "FOXP3"}, + {"source_index": 0, "source_name": "DAPI"}, + ] + + +def test_mixed_native_and_brightfield_requests_fail_but_native_models_can_coexist(): + with pytest.raises(ValueError, match="cannot be combined"): + _reject_mixed_native_and_brightfield_models(["kronos", "uni"]) + _reject_mixed_native_and_brightfield_models(["kronos", "kronos2"]) + + +def test_native_parameter_namespaces_are_finite_and_bound_to_requested_models(): + with pytest.raises(ValueError, match="finite JSON"): + RunConfig( + models=["kronos2"], + native_multiplex_params={"kronos2": {"value": float("nan")}}, + ) + with pytest.raises(ValueError, match="not requested native"): + resolve_multiplex_source_config( + "unused.zarr", + RunConfig( + models=["kronos2"], + native_multiplex_params={"kronos": {"value": 1}}, + ), + ) + + +def test_native_parameters_are_filtered_to_each_geometry_group(): + cfg = RunConfig( + models=["kronos", "kronos2"], + no_seg=True, + native_multiplex_params={"kronos2": {"additional_markers": "panel.csv"}}, + ) + _groups, group_cfgs, _run_hash = resolve_run(cfg) + by_model = {group.models[0]: group for group in group_cfgs} + assert by_model["kronos"].native_multiplex_params == {} + assert by_model["kronos2"].native_multiplex_params == { + "kronos2": {"additional_markers": "panel.csv"} + } diff --git a/tests/test_slide_embedders.py b/tests/test_slide_embedders.py index af00b9a..51f00a7 100644 --- a/tests/test_slide_embedders.py +++ b/tests/test_slide_embedders.py @@ -153,7 +153,7 @@ def test_patch_registry_ignores_slide_encoders_section(): "gigapath", "gigapath_flash", "conch", "conch_v1_5", "h_optimus_0", "h0_mini", "gpfm", "midnight", "openmidnight", "openpath", - "ctranspath", "hibou_l", "hibou_b", "kronos", "phikon", "phikon_v2", + "ctranspath", "hibou_l", "hibou_b", "kronos", "kronos2", "phikon", "phikon_v2", "lunit_dino", "lunit_dino8", "lunit_bt", "lunit_mocov2", "lunit_swav", "sp22m", "retccl", "hipt", "h_optimus_1", "virchow", "musk", "mstar", "kaiko_vitl", "quiltnet", "biomedclip", "plip", "keep", diff --git a/uv.lock b/uv.lock index 0e97eab..3380dac 100644 --- a/uv.lock +++ b/uv.lock @@ -282,6 +282,12 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cb/43/fff37380b704b60fad17c298b631044d98e3db8b9721df7897c967facf51/annsel-0.1.2-py3-none-any.whl", hash = "sha256:9659f60f4ee5545bc030341b6b8e932e0a615b95ac83d9cb6e2bd50951ae26da", size = 19120, upload-time = "2026-01-06T00:07:18.69Z" }, ] +[[package]] +name = "antlr4-python3-runtime" +version = "4.9.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/3e/38/7859ff46355f76f8d19459005ca000b6e7012f2f1ca597746cbcd1fbfe5e/antlr4-python3-runtime-4.9.3.tar.gz", hash = "sha256:f224469b4168294902bb1efa80a8bf7855f24c99aef99cbefc1bcd3cce77881b", size = 117034, upload-time = "2021-11-06T17:52:23.524Z" } + [[package]] name = "array-api-compat" version = "1.15.0" @@ -546,74 +552,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0c/58/bd257695f39d05594ca4ad60df5bcb7e32247f9951fd09a9b8edb82d1daa/contourpy-1.3.3-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:3d1a3799d62d45c18bafd41c5fa05120b96a28079f2393af559b843d1a966a77", size = 225315, upload-time = "2025-07-26T12:02:58.801Z" }, ] -[[package]] -name = "cuda-bindings" -version = "13.3.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "cuda-pathfinder", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/f9/52/50673d25e46d199556f827514bf646a49471d50538c5e577201245b348a9/cuda_bindings-13.3.0-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:244a167d81a9d07d3209d02c8e02e848c2b7c38f4d01e8e4d1f9620b173ae006", size = 6051409, upload-time = "2026-05-27T03:59:01.648Z" }, - { url = "https://files.pythonhosted.org/packages/88/ee/e8f4bdfb808c3689539b7c035d63b6dac9f236b2d6f807f18c7f5f3ef879/cuda_bindings-13.3.0-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:820ab45be3f39a088f39c4a04eb8b852d0b339bff8c518da5c258882b8a4e21b", size = 6671833, upload-time = "2026-05-27T03:59:03.761Z" }, - { url = "https://files.pythonhosted.org/packages/1f/e0/4b3fdba08ff177e9451f376a4ba2df18d76f9158e6a16cdc062bd83db9fa/cuda_bindings-13.3.0-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f72f67f5790e7fa51e3f893e007e49567573ccdf4ed1ca988fbbbd36cd77847c", size = 6020531, upload-time = "2026-05-27T03:59:07.942Z" }, - { url = "https://files.pythonhosted.org/packages/04/40/a2ea4d8f032bfd6c220d50b6f92cd61f33d48f31959da39ed1b178cfee54/cuda_bindings-13.3.0-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a99c3b8d584f266c616bd0f30c7cd83e33553e3ef2abad41ff5a74fbc033a69a", size = 6653764, upload-time = "2026-05-27T03:59:09.981Z" }, - { url = "https://files.pythonhosted.org/packages/ae/a0/156efe7816699c2de1ea2395031db7d010b7af23c243563a3ee6f0ecc1de/cuda_bindings-13.3.0-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a7698fcc4577aa96372866f4d0c9a6cf686cd5c90eab94581c29d37fe6600542", size = 5914803, upload-time = "2026-05-27T03:59:14.011Z" }, - { url = "https://files.pythonhosted.org/packages/51/91/510aae64d53227b5b36db6bfaea41514b66d92cd65ddc43aa49566f18313/cuda_bindings-13.3.0-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:abd908f651160d12c45c5714a38ee102a1173a55433c0d1509ec0e8293beb4a6", size = 6472506, upload-time = "2026-05-27T03:59:16.551Z" }, - { url = "https://files.pythonhosted.org/packages/01/53/2ef49e5b3734a5531b2ba5d726cba724d9cbb262404e586ed61070604826/cuda_bindings-13.3.0-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4a801fa30e75d25b74252123aefc746b6c4275624d2b8640632dd1dfeeaa1f88", size = 6008814, upload-time = "2026-05-27T03:59:20.921Z" }, - { url = "https://files.pythonhosted.org/packages/2f/cb/3a9fcf0651e0a49b4d0f1955837ce079245b27086c22fb2f253039bdf324/cuda_bindings-13.3.0-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:94d40ef7b4bdd9dce0244a1baa132e0e538f1eb2c0d162fb3648a15e48515365", size = 6531477, upload-time = "2026-05-27T03:59:23.391Z" }, - { url = "https://files.pythonhosted.org/packages/2b/0f/6987c5ee98f117317a85650ddc79480a3fa59a573ae1c923d0722b56ae71/cuda_bindings-13.3.0-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e5911bea15810b749a8077f8c45423ed785d51618b8e8664dea1fc8f5a2a76c8", size = 5807073, upload-time = "2026-05-27T03:59:28.218Z" }, - { url = "https://files.pythonhosted.org/packages/f6/ab/46ceee07dc19f18a5d1c28d592750ed9dbdc803077eb083576a442c9938c/cuda_bindings-13.3.0-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e2870fed7707a37f8af0c02364b05f355ebe8921604e8c68eb56cf66867e0798", size = 6354325, upload-time = "2026-05-27T03:59:30.715Z" }, -] - -[[package]] -name = "cuda-pathfinder" -version = "1.5.5" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/11/c8/26f2e4aae92f11522a96043892ba39a90eac610d5242523aa863212bc1c7/cuda_pathfinder-1.5.5-py3-none-any.whl", hash = "sha256:0228c023f95d1480f143ef5c8922d27a2ab052087a942e81dc289c9eb8f91689", size = 51671, upload-time = "2026-05-27T01:21:25.413Z" }, -] - -[[package]] -name = "cuda-toolkit" -version = "13.0.2" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/57/b2/453099f5f3b698d7d0eab38916aac44c7f76229f451709e2eb9db6615dcd/cuda_toolkit-13.0.2-py2.py3-none-any.whl", hash = "sha256:b198824cf2f54003f50d64ada3a0f184b42ca0846c1c94192fa269ecd97a66eb", size = 2364, upload-time = "2025-12-19T23:24:07.328Z" }, -] - -[package.optional-dependencies] -cudart = [ - { name = "nvidia-cuda-runtime", marker = "sys_platform == 'linux'" }, -] -cufft = [ - { name = "nvidia-cufft", marker = "sys_platform == 'linux'" }, -] -cufile = [ - { name = "nvidia-cufile", marker = "sys_platform == 'linux'" }, -] -cupti = [ - { name = "nvidia-cuda-cupti", marker = "sys_platform == 'linux'" }, -] -curand = [ - { name = "nvidia-curand", marker = "sys_platform == 'linux'" }, -] -cusolver = [ - { name = "nvidia-cusolver", marker = "sys_platform == 'linux'" }, -] -cusparse = [ - { name = "nvidia-cusparse", marker = "sys_platform == 'linux'" }, -] -nvjitlink = [ - { name = "nvidia-nvjitlink", marker = "sys_platform == 'linux'" }, -] -nvrtc = [ - { name = "nvidia-cuda-nvrtc", marker = "sys_platform == 'linux'" }, -] -nvtx = [ - { name = "nvidia-nvtx", marker = "sys_platform == 'linux'" }, -] - [[package]] name = "cycler" version = "0.12.1" @@ -2122,119 +2060,97 @@ wheels = [ ] [[package]] -name = "nvidia-cublas" -version = "13.1.1.3" +name = "nvidia-cublas-cu12" +version = "12.4.5.8" source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "nvidia-cuda-nvrtc", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, -] wheels = [ - { url = "https://files.pythonhosted.org/packages/a7/a1/0bd24ee8c8d03adac032fd2909426a00c88f8c57961b1277ded97f91119f/nvidia_cublas-13.1.1.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:b7a210458267ac818974c53038fbec2e969d5c99f305ab15c72522fa9f001dd5", size = 542848918, upload-time = "2026-04-08T18:46:22.985Z" }, - { url = "https://files.pythonhosted.org/packages/3b/cd/154ca20c38269e05eff77c1464e6c1da89f50a6390b565e9d82e06bc11e1/nvidia_cublas-13.1.1.3-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:37936a16db8fe4ac1f065c2139360608a543a09275cb1a1af612e08cfa065436", size = 423138758, upload-time = "2026-04-08T18:46:58.655Z" }, + { url = "https://files.pythonhosted.org/packages/ae/71/1c91302526c45ab494c23f61c7a84aa568b8c1f9d196efa5993957faf906/nvidia_cublas_cu12-12.4.5.8-py3-none-manylinux2014_x86_64.whl", hash = "sha256:2fc8da60df463fdefa81e323eef2e36489e1c94335b5358bcb38360adf75ac9b", size = 363438805, upload-time = "2024-04-03T20:57:06.025Z" }, ] [[package]] -name = "nvidia-cuda-cupti" -version = "13.0.85" +name = "nvidia-cuda-cupti-cu12" +version = "12.4.127" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/2a/2a/80353b103fc20ce05ef51e928daed4b6015db4aaa9162ed0997090fe2250/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_aarch64.whl", hash = "sha256:796bd679890ee55fb14a94629b698b6db54bcfd833d391d5e94017dd9d7d3151", size = 10310827, upload-time = "2025-09-04T08:26:42.012Z" }, - { url = "https://files.pythonhosted.org/packages/33/6d/737d164b4837a9bbd202f5ae3078975f0525a55730fe871d8ed4e3b952b0/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_x86_64.whl", hash = "sha256:4eb01c08e859bf924d222250d2e8f8b8ff6d3db4721288cf35d14252a4d933c8", size = 10715597, upload-time = "2025-09-04T08:26:51.312Z" }, + { url = "https://files.pythonhosted.org/packages/67/42/f4f60238e8194a3106d06a058d494b18e006c10bb2b915655bd9f6ea4cb1/nvidia_cuda_cupti_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:9dec60f5ac126f7bb551c055072b69d85392b13311fcc1bcda2202d172df30fb", size = 13813957, upload-time = "2024-04-03T20:55:01.564Z" }, ] [[package]] -name = "nvidia-cuda-nvrtc" -version = "13.0.88" +name = "nvidia-cuda-nvrtc-cu12" +version = "12.4.127" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/c3/68/483a78f5e8f31b08fb1bb671559968c0ca3a065ac7acabfc7cee55214fd6/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:ad9b6d2ead2435f11cbb6868809d2adeeee302e9bb94bcf0539c7a40d80e8575", size = 90215200, upload-time = "2025-09-04T08:28:44.204Z" }, - { url = "https://files.pythonhosted.org/packages/b7/dc/6bb80850e0b7edd6588d560758f17e0550893a1feaf436807d64d2da040f/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d27f20a0ca67a4bb34268a5e951033496c5b74870b868bacd046b1b8e0c3267b", size = 43015449, upload-time = "2025-09-04T08:28:20.239Z" }, + { url = "https://files.pythonhosted.org/packages/2c/14/91ae57cd4db3f9ef7aa99f4019cfa8d54cb4caa7e00975df6467e9725a9f/nvidia_cuda_nvrtc_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:a178759ebb095827bd30ef56598ec182b85547f1508941a3d560eb7ea1fbf338", size = 24640306, upload-time = "2024-04-03T20:56:01.463Z" }, ] [[package]] -name = "nvidia-cuda-runtime" -version = "13.0.96" +name = "nvidia-cuda-runtime-cu12" +version = "12.4.127" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/87/4f/17d7b9b8e285199c58ce28e31b5c5bbaa4d8271af06a89b6405258245de2/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ef9bcbe90493a2b9d810e43d249adb3d02e98dd30200d86607d8d02687c43f55", size = 2261060, upload-time = "2025-10-09T08:55:15.78Z" }, - { url = "https://files.pythonhosted.org/packages/2e/24/d1558f3b68b1d26e706813b1d10aa1d785e4698c425af8db8edc3dced472/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7f82250d7782aa23b6cfe765ecc7db554bd3c2870c43f3d1821f1d18aebf0548", size = 2243632, upload-time = "2025-10-09T08:55:36.117Z" }, + { url = "https://files.pythonhosted.org/packages/ea/27/1795d86fe88ef397885f2e580ac37628ed058a92ed2c39dc8eac3adf0619/nvidia_cuda_runtime_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:64403288fa2136ee8e467cdc9c9427e0434110899d07c779f25b5c068934faa5", size = 883737, upload-time = "2024-04-03T20:54:51.355Z" }, ] [[package]] -name = "nvidia-cudnn-cu13" -version = "9.20.0.48" +name = "nvidia-cudnn-cu12" +version = "9.1.0.70" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-cublas-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/56/c5/83384d846b2fd17c44bd499b36c75a45ed4f095fbbb2252294e89cea5c5c/nvidia_cudnn_cu13-9.20.0.48-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:e31454ae00094b0c55319d9d15b6fa2fc50a9e1c0f5c8c80fb75258234e731e1", size = 444574296, upload-time = "2026-03-09T19:28:27.751Z" }, - { url = "https://files.pythonhosted.org/packages/6e/5e/edb9c0ae051602c3ccaffe424256463636d639e27d7f302dde9975ef9e7a/nvidia_cudnn_cu13-9.20.0.48-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:0c45dd8eeb50b603f07995b1b300c62ffe6a1980482b82b3bcf94a4ca9d49304", size = 366173588, upload-time = "2026-03-09T19:29:34.474Z" }, + { url = "https://files.pythonhosted.org/packages/9f/fd/713452cd72343f682b1c7b9321e23829f00b842ceaedcda96e742ea0b0b3/nvidia_cudnn_cu12-9.1.0.70-py3-none-manylinux2014_x86_64.whl", hash = "sha256:165764f44ef8c61fcdfdfdbe769d687e06374059fbb388b6c89ecb0e28793a6f", size = 664752741, upload-time = "2024-04-22T15:24:15.253Z" }, ] [[package]] -name = "nvidia-cufft" -version = "12.0.0.61" +name = "nvidia-cufft-cu12" +version = "11.2.1.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" }, - { url = "https://files.pythonhosted.org/packages/a8/2f/7b57e29836ea8714f81e9898409196f47d772d5ddedddf1592eadb8ab743/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6c44f692dce8fd5ffd3e3df134b6cdb9c2f72d99cf40b62c32dde45eea9ddad3", size = 214085489, upload-time = "2025-09-04T08:31:56.044Z" }, + { name = "nvidia-nvjitlink-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] - -[[package]] -name = "nvidia-cufile" -version = "1.15.1.6" -source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/3f/70/4f193de89a48b71714e74602ee14d04e4019ad36a5a9f20c425776e72cd6/nvidia_cufile-1.15.1.6-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:08a3ecefae5a01c7f5117351c64f17c7c62efa5fffdbe24fc7d298da19cd0b44", size = 1223672, upload-time = "2025-09-04T08:32:22.779Z" }, - { url = "https://files.pythonhosted.org/packages/ab/73/cc4a14c9813a8a0d509417cf5f4bdaba76e924d58beb9864f5a7baceefbf/nvidia_cufile-1.15.1.6-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:bdc0deedc61f548bddf7733bdc216456c2fdb101d020e1ab4b88d232d5e2f6d1", size = 1136992, upload-time = "2025-09-04T08:32:14.119Z" }, + { url = "https://files.pythonhosted.org/packages/27/94/3266821f65b92b3138631e9c8e7fe1fb513804ac934485a8d05776e1dd43/nvidia_cufft_cu12-11.2.1.3-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f083fc24912aa410be21fa16d157fed2055dab1cc4b6934a0e03cba69eb242b9", size = 211459117, upload-time = "2024-04-03T20:57:40.402Z" }, ] [[package]] -name = "nvidia-curand" -version = "10.4.0.35" +name = "nvidia-curand-cu12" +version = "10.3.5.147" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/1e/72/7c2ae24fb6b63a32e6ae5d241cc65263ea18d08802aaae087d9f013335a2/nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:133df5a7509c3e292aaa2b477afd0194f06ce4ea24d714d616ff36439cee349a", size = 61962106, upload-time = "2025-08-04T10:21:41.128Z" }, - { url = "https://files.pythonhosted.org/packages/a5/9f/be0a41ca4a4917abf5cb9ae0daff1a6060cc5de950aec0396de9f3b52bc5/nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:1aee33a5da6e1db083fe2b90082def8915f30f3248d5896bcec36a579d941bfc", size = 59544258, upload-time = "2025-08-04T10:22:03.992Z" }, + { url = "https://files.pythonhosted.org/packages/8a/6d/44ad094874c6f1b9c654f8ed939590bdc408349f137f9b98a3a23ccec411/nvidia_curand_cu12-10.3.5.147-py3-none-manylinux2014_x86_64.whl", hash = "sha256:a88f583d4e0bb643c49743469964103aa59f7f708d862c3ddb0fc07f851e3b8b", size = 56305206, upload-time = "2024-04-03T20:58:08.722Z" }, ] [[package]] -name = "nvidia-cusolver" -version = "12.0.4.66" +name = "nvidia-cusolver-cu12" +version = "11.6.1.9" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "nvidia-cusparse", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "nvidia-nvjitlink", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-cublas-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-cusparse-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-nvjitlink-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" }, - { url = "https://files.pythonhosted.org/packages/5f/67/cba3777620cdacb99102da4042883709c41c709f4b6323c10781a9c3aa34/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:0a759da5dea5c0ea10fd307de75cdeb59e7ea4fcb8add0924859b944babf1112", size = 200941980, upload-time = "2025-09-04T08:33:22.767Z" }, + { url = "https://files.pythonhosted.org/packages/3a/e1/5b9089a4b2a4790dfdea8b3a006052cfecff58139d5a4e34cb1a51df8d6f/nvidia_cusolver_cu12-11.6.1.9-py3-none-manylinux2014_x86_64.whl", hash = "sha256:19e33fa442bcfd085b3086c4ebf7e8debc07cfe01e11513cc6d332fd918ac260", size = 127936057, upload-time = "2024-04-03T20:58:28.735Z" }, ] [[package]] -name = "nvidia-cusparse" -version = "12.6.3.3" +name = "nvidia-cusparse-cu12" +version = "12.3.1.170" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-nvjitlink-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" }, - { url = "https://files.pythonhosted.org/packages/fa/18/623c77619c31d62efd55302939756966f3ecc8d724a14dab2b75f1508850/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:2b3c89c88d01ee0e477cb7f82ef60a11a4bcd57b6b87c33f789350b59759360b", size = 145942937, upload-time = "2025-09-04T08:33:58.029Z" }, + { url = "https://files.pythonhosted.org/packages/db/f7/97a9ea26ed4bbbfc2d470994b8b4f338ef663be97b8f677519ac195e113d/nvidia_cusparse_cu12-12.3.1.170-py3-none-manylinux2014_x86_64.whl", hash = "sha256:ea4f11a2904e2a8dc4b1833cc1b5181cde564edd0d5cd33e3c168eff2d1863f1", size = 207454763, upload-time = "2024-04-03T20:58:59.995Z" }, ] [[package]] -name = "nvidia-cusparselt-cu13" -version = "0.8.1" +name = "nvidia-cusparselt-cu12" +version = "0.6.2" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/46/e1/cdc1797eadf82d3a9a575a19b33fdc871a97edbec42c00b5b5e914f4aff4/nvidia_cusparselt_cu13-0.8.1-py3-none-manylinux2014_aarch64.whl", hash = "sha256:4dca476c50bf4780d46cd0bfbd82e2bc10a08e4fef7950917ce8d7578d22a23f", size = 221051344, upload-time = "2025-09-05T18:49:51.289Z" }, - { url = "https://files.pythonhosted.org/packages/34/7d/2661f2fb3ac4302f3a246f5fc030213ac60c1fe0bce84f9783dbd831dbb7/nvidia_cusparselt_cu13-0.8.1-py3-none-manylinux2014_x86_64.whl", hash = "sha256:786ce87568c303fadb5afcc7102d454cd3040d75f6f8626f5db460d1871f4dd0", size = 170148586, upload-time = "2025-09-05T18:50:50.248Z" }, + { url = "https://files.pythonhosted.org/packages/78/a8/bcbb63b53a4b1234feeafb65544ee55495e1bb37ec31b999b963cbccfd1d/nvidia_cusparselt_cu12-0.6.2-py3-none-manylinux2014_x86_64.whl", hash = "sha256:df2c24502fd76ebafe7457dbc4716b2fec071aabaed4fb7691a201cde03704d9", size = 150057751, upload-time = "2024-07-23T02:35:53.074Z" }, ] [[package]] @@ -2247,39 +2163,27 @@ wheels = [ ] [[package]] -name = "nvidia-nccl-cu13" -version = "2.29.7" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/72/0d/daf50d44177ee0cbc7ff0a0c91eb5ff676c82be42f9a970bc7597f440c3a/nvidia_nccl_cu13-2.29.7-py3-none-manylinux_2_18_aarch64.whl", hash = "sha256:674a12383e3c38a1bcccae7d4f3633b37852230b6047883cb2f4c2d1b36d9bf5", size = 206014712, upload-time = "2026-03-03T05:34:20.843Z" }, - { url = "https://files.pythonhosted.org/packages/67/f4/58e4e91b6919367c7aafb8e36fce9aad1a3047e536bf7e2fd560927d3a4c/nvidia_nccl_cu13-2.29.7-py3-none-manylinux_2_18_x86_64.whl", hash = "sha256:edd81538446786ec3b73972543e53bb43bcaf0bfc8ef76cb679fcc390ffe136d", size = 205976000, upload-time = "2026-03-03T05:36:24.472Z" }, -] - -[[package]] -name = "nvidia-nvjitlink" -version = "13.0.88" +name = "nvidia-nccl-cu12" +version = "2.21.5" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/56/7a/123e033aaff487c77107195fa5a2b8686795ca537935a24efae476c41f05/nvidia_nvjitlink-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:13a74f429e23b921c1109976abefacc69835f2f433ebd323d3946e11d804e47b", size = 40713933, upload-time = "2025-09-04T08:35:43.553Z" }, - { url = "https://files.pythonhosted.org/packages/ab/2c/93c5250e64df4f894f1cbb397c6fd71f79813f9fd79d7cd61de3f97b3c2d/nvidia_nvjitlink-13.0.88-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e931536ccc7d467a98ba1d8b89ff7fa7f1fa3b13f2b0069118cd7f47bff07d0c", size = 38768748, upload-time = "2025-09-04T08:35:20.008Z" }, + { url = "https://files.pythonhosted.org/packages/df/99/12cd266d6233f47d00daf3a72739872bdc10267d0383508b0b9c84a18bb6/nvidia_nccl_cu12-2.21.5-py3-none-manylinux2014_x86_64.whl", hash = "sha256:8579076d30a8c24988834445f8d633c697d42397e92ffc3f63fa26766d25e0a0", size = 188654414, upload-time = "2024-04-03T15:32:57.427Z" }, ] [[package]] -name = "nvidia-nvshmem-cu13" -version = "3.4.5" +name = "nvidia-nvjitlink-cu12" +version = "12.4.127" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/dc/0f/05cc9c720236dcd2db9c1ab97fff629e96821be2e63103569da0c9b72f19/nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6dc2a197f38e5d0376ad52cd1a2a3617d3cdc150fd5966f4aee9bcebb1d68fe9", size = 60215947, upload-time = "2025-09-06T00:32:20.022Z" }, - { url = "https://files.pythonhosted.org/packages/3c/35/a9bf80a609e74e3b000fef598933235c908fcefcef9026042b8e6dfde2a9/nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:290f0a2ee94c9f3687a02502f3b9299a9f9fe826e6d0287ee18482e78d495b80", size = 60412546, upload-time = "2025-09-06T00:32:41.564Z" }, + { url = "https://files.pythonhosted.org/packages/ff/ff/847841bacfbefc97a00036e0fce5a0f086b640756dc38caea5e1bb002655/nvidia_nvjitlink_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:06b3b9b25bf3f8af351d664978ca26a16d2c5127dbd53c0497e28d1fb9611d57", size = 21066810, upload-time = "2024-04-03T20:59:46.957Z" }, ] [[package]] -name = "nvidia-nvtx" -version = "13.0.85" +name = "nvidia-nvtx-cu12" +version = "12.4.127" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/c2/f3/d86c845465a2723ad7e1e5c36dcd75ddb82898b3f53be47ebd429fb2fa5d/nvidia_nvtx-13.0.85-py3-none-manylinux1_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:4936d1d6780fbe68db454f5e72a42ff64d1fd6397df9f363ae786930fd5c1cd4", size = 148047, upload-time = "2025-09-04T08:29:01.761Z" }, - { url = "https://files.pythonhosted.org/packages/a8/64/3708a90d1ebe202ffdeb7185f878a3c84d15c2b2c31858da2ce0583e2def/nvidia_nvtx-13.0.85-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cb7780edb6b14107373c835bf8b72e7a178bac7367e23da7acb108f973f157a6", size = 148878, upload-time = "2025-09-04T08:28:53.627Z" }, + { url = "https://files.pythonhosted.org/packages/87/20/199b8713428322a2f22b722c62b8cc278cc53dffa9705d744484b5035ee9/nvidia_nvtx_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:781e950d9b9f60d8241ccea575b32f5105a5baf4c2351cab5256a24869f12a1a", size = 99144, upload-time = "2024-04-03T20:56:12.406Z" }, ] [[package]] @@ -2358,6 +2262,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5b/fc/7650df09dce0edfbaa1f2898c5eb0e0523ae5964dadb83855beb52166505/ome_zarr_models-1.6-py3-none-any.whl", hash = "sha256:d855a6fe885aa2613f64cd124ff8453d75854531bb3282eeaa916e976c7bcf66", size = 64251, upload-time = "2026-03-10T10:14:27.177Z" }, ] +[[package]] +name = "omegaconf" +version = "2.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "antlr4-python3-runtime" }, + { name = "pyyaml" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ce/3d/e4b57b8d9008c6ebe0d5eff901f91d5700cf7bdb8c8863df817463a7fd5e/omegaconf-2.3.1.tar.gz", hash = "sha256:e5e7de64aeebeddaf8e6d3f7a783b32ac2a01c0fbd9c878012caecb891a1f42a", size = 3298472, upload-time = "2026-06-11T05:05:12.885Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a4/0e/152509871bf30df6fc38569f52a2db9b55dd41aae957adae50a053ac7778/omegaconf-2.3.1-py3-none-any.whl", hash = "sha256:3d701d14e9a8828f1edd28bb70b725908b34277cdd72cf7d6a83f94dadc6b6a0", size = 79502, upload-time = "2026-06-11T05:05:09.954Z" }, +] + [[package]] name = "open-clip-torch" version = "3.3.0" @@ -3371,7 +3288,7 @@ wheels = [ [[package]] name = "raw2features" -version = "0.2.0" +version = "0.2.1" source = { editable = "." } dependencies = [ { name = "jsonschema" }, @@ -3470,6 +3387,17 @@ kronos = [ { name = "torchvision" }, { name = "transformers" }, ] +kronos2 = [ + { name = "huggingface-hub" }, + { name = "omegaconf" }, + { name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.14' and sys_platform != 'win32') or (python_full_version > '3.11' and sys_platform == 'win32')" }, + { name = "pandas", version = "3.0.3", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.14' and sys_platform != 'win32') or (python_full_version <= '3.11' and sys_platform == 'win32')" }, + { name = "safetensors" }, + { name = "timm" }, + { name = "torch" }, + { name = "transformers" }, + { name = "xformers", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] madeleine = [ { name = "einops" }, { name = "einops-exts" }, @@ -3602,6 +3530,7 @@ requires-dist = [ { name = "gdown", marker = "extra == 'tangle'", specifier = ">=5.2" }, { name = "geopandas", marker = "extra == 'spatialdata'", specifier = ">=0.14" }, { name = "h5py", marker = "extra == 'h5'", specifier = ">=3.0" }, + { name = "huggingface-hub", marker = "extra == 'kronos2'", specifier = ">=0.28.1,<1.0" }, { name = "huggingface-hub", marker = "extra == 'models'", specifier = ">=0.24" }, { name = "imagecodecs", marker = "extra == 'kronos'" }, { name = "jsonschema", specifier = ">=4.18" }, @@ -3609,9 +3538,11 @@ requires-dist = [ { name = "numcodecs", marker = "extra == 'zarr'", specifier = ">=0.15" }, { name = "numpy", specifier = ">=2.0" }, { name = "ome-zarr", marker = "extra == 'zarr'", specifier = ">=0.10" }, + { name = "omegaconf", marker = "extra == 'kronos2'", specifier = ">=2.3" }, { name = "open-clip-torch", marker = "extra == 'musk'" }, { name = "open-clip-torch", marker = "extra == 'open-clip'", specifier = ">=3.0" }, { name = "opencv-python-headless", marker = "extra == 'image'", specifier = ">=4.9" }, + { name = "pandas", marker = "extra == 'kronos2'", specifier = ">=2.2.2" }, { name = "pandas", marker = "extra == 'tables'", specifier = ">=2.2" }, { name = "peft", marker = "extra == 'seal'" }, { name = "pillow", marker = "extra == 'image'", specifier = ">=11.0" }, @@ -3634,6 +3565,7 @@ requires-dist = [ { name = "raw2features", extras = ["zarr", "image", "torch", "models", "stain"], marker = "extra == 'all'" }, { name = "s3fs", marker = "extra == 's3'", specifier = ">=2024.0" }, { name = "sacremoses", marker = "extra == 'prism'" }, + { name = "safetensors", marker = "extra == 'kronos2'", specifier = ">=0.4" }, { name = "safetensors", marker = "extra == 'models'", specifier = ">=0.4" }, { name = "scanpy", marker = "extra == 'seal'" }, { name = "scikit-learn", marker = "extra == 'stain'", specifier = ">=1.3" }, @@ -3642,16 +3574,20 @@ requires-dist = [ { name = "spatialdata", marker = "extra == 'spatialdata'", specifier = ">=0.2" }, { name = "tifffile", marker = "extra == 'kronos'", specifier = ">=2024" }, { name = "timm", marker = "extra == 'grandqc'", specifier = ">=1.0" }, + { name = "timm", marker = "extra == 'kronos2'", specifier = "==1.0.19" }, { name = "timm", marker = "extra == 'models'", specifier = ">=1.0" }, + { name = "torch", marker = "extra == 'kronos2'", specifier = "==2.6.0" }, { name = "torch", marker = "extra == 'torch'", specifier = ">=2.2" }, { name = "torchvision", marker = "extra == 'torch'", specifier = ">=0.17" }, { name = "tqdm", specifier = ">=4.66" }, { name = "transformers", marker = "extra == 'hibou'", specifier = "<5" }, + { name = "transformers", marker = "extra == 'kronos2'", specifier = "==4.56.0" }, { name = "transformers", marker = "extra == 'models'", specifier = ">=4.40" }, { name = "typer", specifier = ">=0.12" }, + { name = "xformers", marker = "platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'kronos2'", specifier = "==0.0.29.post3" }, { name = "zarr", marker = "extra == 'zarr'", specifier = ">=3.1" }, ] -provides-extras = ["all", "zarr", "image", "stain", "torch", "models", "conch", "hibou", "kronos", "musk", "open-clip", "grandqc", "prism", "madeleine", "gigapath-slide", "chief", "tangle", "seal", "spatialdata", "h5", "tables", "s3", "benchmark"] +provides-extras = ["all", "zarr", "image", "stain", "torch", "models", "conch", "hibou", "kronos", "kronos2", "musk", "open-clip", "grandqc", "prism", "madeleine", "gigapath-slide", "chief", "tangle", "seal", "spatialdata", "h5", "tables", "s3", "benchmark"] [package.metadata.requires-dev] dev = [ @@ -4657,14 +4593,14 @@ wheels = [ [[package]] name = "sympy" -version = "1.14.0" +version = "1.13.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "mpmath" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/83/d3/803453b36afefb7c2bb238361cd4ae6125a569b4db67cd9e79846ba2d68c/sympy-1.14.0.tar.gz", hash = "sha256:d3d3fe8df1e5a0b42f0e7bdf50541697dbe7d23746e894990c030e2b05e72517", size = 7793921, upload-time = "2025-04-27T18:05:01.611Z" } +sdist = { url = "https://files.pythonhosted.org/packages/ca/99/5a5b6f19ff9f083671ddf7b9632028436167cd3d33e11015754e41b249a4/sympy-1.13.1.tar.gz", hash = "sha256:9cebf7e04ff162015ce31c9c6c9144daa34a93bd082f54fd8f12deca4f47515f", size = 7533040, upload-time = "2024-07-19T09:26:51.238Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, + { url = "https://files.pythonhosted.org/packages/b2/fe/81695a1aa331a842b582453b605175f419fe8540355886031328089d840a/sympy-1.13.1-py3-none-any.whl", hash = "sha256:db36cdc64bf61b9b24578b6f7bab1ecdd2452cf008f34faa33776680c26d66f8", size = 6189177, upload-time = "2024-07-19T09:26:48.863Z" }, ] [[package]] @@ -4727,7 +4663,7 @@ wheels = [ [[package]] name = "timm" -version = "1.0.27" +version = "1.0.19" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "huggingface-hub" }, @@ -4736,9 +4672,9 @@ dependencies = [ { name = "torch" }, { name = "torchvision" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/08/54/ece85b0eef3700c90db8271a43669b05a0ebbe2edb1962329c34374a297e/timm-1.0.27.tar.gz", hash = "sha256:315dfe63186ca9fb7ff941268941231fd5be259f2b4bb4afa28560ae1015cb9a", size = 2439861, upload-time = "2026-05-08T19:38:36.844Z" } +sdist = { url = "https://files.pythonhosted.org/packages/1c/78/0789838cf20ba1cc09907914a008c1823d087132b48aa1ccde5e7934175a/timm-1.0.19.tar.gz", hash = "sha256:6e71e1f67ac80c229d3a78ca58347090514c508aeba8f2e2eb5289eda86e9f43", size = 2353261, upload-time = "2025-07-24T03:04:05.281Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/1f/2e/26bab7686ff4aed48f8f5f6c23e2aa37b7a37ddd9effe3aa61e908fd518f/timm-1.0.27-py3-none-any.whl", hash = "sha256:5ff07c9ddf53cbada88eab1c93ff175c64cab683b5a2fddf863bcee985926f89", size = 2589280, upload-time = "2026-05-08T19:38:35.034Z" }, + { url = "https://files.pythonhosted.org/packages/dc/74/661c63260cccf19ed5932e8b3f22f95ecd8bb34b9d9e6af9e1e7b961f254/timm-1.0.19-py3-none-any.whl", hash = "sha256:c07b56c32f3d3226c656f75c1b5479c08eb34eefed927c82fd8751a852f47931", size = 2497950, upload-time = "2025-07-24T03:04:03.097Z" }, ] [[package]] @@ -4778,55 +4714,49 @@ wheels = [ [[package]] name = "torch" -version = "2.12.0" +version = "2.6.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cuda-bindings", marker = "sys_platform == 'linux'" }, - { name = "cuda-toolkit", extra = ["cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "sys_platform == 'linux'" }, { name = "filelock" }, { name = "fsspec" }, { name = "jinja2" }, { name = "networkx" }, - { name = "nvidia-cublas", marker = "sys_platform == 'linux'" }, - { name = "nvidia-cudnn-cu13", marker = "sys_platform == 'linux'" }, - { name = "nvidia-cusparselt-cu13", marker = "sys_platform == 'linux'" }, - { name = "nvidia-nccl-cu13", marker = "sys_platform == 'linux'" }, - { name = "nvidia-nvshmem-cu13", marker = "sys_platform == 'linux'" }, - { name = "setuptools" }, + { name = "nvidia-cublas-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-cupti-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-nvrtc-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-runtime-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cudnn-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cufft-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-curand-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusolver-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusparse-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusparselt-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nccl-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nvtx-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "setuptools", marker = "python_full_version >= '3.12'" }, { name = "sympy" }, - { name = "triton", marker = "sys_platform == 'linux'" }, + { name = "triton", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "typing-extensions" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/18/62/131124fb95df03811b8260d1d43dcc5ee85ea1a344b964613d7efe77fb08/torch-2.12.0-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:10802fd383bbfed646212e765a72c37d2185205d4f26eb197a254e8ac7ddcb25", size = 87990344, upload-time = "2026-05-13T14:55:42.154Z" }, - { url = "https://files.pythonhosted.org/packages/12/9c/dda0dbd547dc549839824135f223792fd0e725f28ed0715dda366b7acaa2/torch-2.12.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:c12592630aef72feaf18bd3f197ef587bbfa21131b31c38b23ab2e55fce92e36", size = 426362932, upload-time = "2026-05-13T14:54:15.295Z" }, - { url = "https://files.pythonhosted.org/packages/e2/d2/a7dd5a3f9bdaa7842124e8e2359202b317c48d47d2fc5816fafdf2049adb/torch-2.12.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:415c1b8d0412f67551c8e89a2daca0fb3e56694af0281ba155eaa9da481f58b4", size = 532170085, upload-time = "2026-05-13T14:55:20.788Z" }, - { url = "https://files.pythonhosted.org/packages/12/1b/a61ce2004f9ab0ea8964a6e6168133a127795667639e2ff4f8f2bdb16a65/torch-2.12.0-cp311-cp311-win_amd64.whl", hash = "sha256:dd37188ea325042cb1f6cafa56822b11ada2520c04791a52629b0af25bdfbfd9", size = 122953128, upload-time = "2026-05-13T14:54:52.744Z" }, - { url = "https://files.pythonhosted.org/packages/ef/bb/285d643f254731294c9b595a007eac39db4600a98682d7bca688f42ca164/torch-2.12.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:b41339df93d491435e790ff8bcbae1c0ce777175889bfd1281d119862793e6a2", size = 88010197, upload-time = "2026-05-13T14:55:35.414Z" }, - { url = "https://files.pythonhosted.org/packages/79/81/76debf1db1343bd929bbb5d74c89fb437c2ed88eb144712557e7bd3eea45/torch-2.12.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:8fbef9f108a863e7722a73740998967e3b074742a834fc5be3a535a2befa7057", size = 426376751, upload-time = "2026-05-13T14:55:03.353Z" }, - { url = "https://files.pythonhosted.org/packages/de/f0/80026028b603c4650ff270fc3785bdef4bd6738765a9cc5a0f5a637d65a2/torch-2.12.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:4b4f64c2c2b11f7510d93dd6412b87025ff6eddd6bb61c3b5a3d892ea20c4756", size = 532261691, upload-time = "2026-05-13T14:52:54.453Z" }, - { url = "https://files.pythonhosted.org/packages/b9/c2/64b06cbb7830fb3cd9be13e1158b31a3f36b68e6a209105ee3c9d9480be0/torch-2.12.0-cp312-cp312-win_amd64.whl", hash = "sha256:8b958caff4a14d3a3b0b2dfc6a378f64dda9728a9dad28c08a0db9ce4dafb549", size = 122988114, upload-time = "2026-05-13T14:54:42.153Z" }, - { url = "https://files.pythonhosted.org/packages/86/ca/01896c80ba921676aa45886b2c5b8d774912de2a1f719de48169c6f755cd/torch-2.12.0-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:90dd587a5f61bfe1307148b581e2084fc5bc4a06e2b90a20e9a36b81087ff16b", size = 88009511, upload-time = "2026-05-13T14:54:47.411Z" }, - { url = "https://files.pythonhosted.org/packages/a5/04/52bdaf4787eab6ac7d7f5851dff934e4def0bc8ead9c8fd2b69b3e529699/torch-2.12.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:864392c73b7654f4d2b3ae712f607937d0dbb1101c4555fbb41848106b297f39", size = 426383231, upload-time = "2026-05-13T14:53:32.129Z" }, - { url = "https://files.pythonhosted.org/packages/49/8a/94bdecd13f5aaa90d45920b89789d9fe7c6f4af8c3cdd7ce01fcb59908fc/torch-2.12.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:5d6b560dfa7d56291c07d615c3bb73e8d9943d9b6d87f76cd0d9d570c4797fa6", size = 532269288, upload-time = "2026-05-13T14:53:49.423Z" }, - { url = "https://files.pythonhosted.org/packages/3e/2f/bdbaaa267de519ef1b73054bf590d8c93c37a266c9a4e24a01bd38b6918f/torch-2.12.0-cp313-cp313-win_amd64.whl", hash = "sha256:3fee918902090ade827643e758e98363278815de583c75d111fdd665ebffde9f", size = 122987706, upload-time = "2026-05-13T14:54:00.335Z" }, - { url = "https://files.pythonhosted.org/packages/9b/ad/e95e822f3538171e22640a7fbe839a1fdb666600bf6487025de2ff03b11a/torch-2.12.0-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:10ee1448a9f304d3b987eb4656f664ba6e4d7b410ca7a5a7c642199777a2cf88", size = 88319556, upload-time = "2026-05-13T14:54:05.574Z" }, - { url = "https://files.pythonhosted.org/packages/b7/07/055d06d985b445d67422d25b033c11cf55bbb81785d4c4e68e28bca5820e/torch-2.12.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:af68dbf403439cae9ceaeaaf92f8352b460787dcd27b92aa05c40dd4a19c0f1e", size = 426397656, upload-time = "2026-05-13T14:52:38.84Z" }, - { url = "https://files.pythonhosted.org/packages/43/94/b0b4fdc3014122e0a7302fb90086d352aa48f2576f0b252561ebb38c01a8/torch-2.12.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:a6a2eebb237d3b1d9ad3b378e86d9b9e0782afdea8b1e0eba6a13646b9b49c07", size = 532183124, upload-time = "2026-05-13T14:53:16.178Z" }, - { url = "https://files.pythonhosted.org/packages/d8/c8/052405e6ad05d3237bfe5a4df78f917773956f8e17813a2d44c059068b74/torch-2.12.0-cp313-cp313t-win_amd64.whl", hash = "sha256:2140e373e9a51a3e22ef62e8d14366d0b470d18f0adf19fdc757368077133a34", size = 123232462, upload-time = "2026-05-13T14:52:27.26Z" }, - { url = "https://files.pythonhosted.org/packages/67/dc/ac069f8d6e8be701535921141055293b0d4819d3d7f224a4612cf157c7f9/torch-2.12.0-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:f7dfae4a519197dfa050e98d8e36378a0fb5899625a875c2b54445005a2e404e", size = 88027282, upload-time = "2026-05-13T14:53:05.258Z" }, - { url = "https://files.pythonhosted.org/packages/33/c3/1c1eb00e34555b536dddf792676026a988d710ed36981aa00499b36b0620/torch-2.12.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:891c769072637c74e9a5a77a3bc782894696d8ffec83b938df8536dee7f0ba78", size = 426386961, upload-time = "2026-05-13T14:51:28.406Z" }, - { url = "https://files.pythonhosted.org/packages/cd/d4/7e730dba0c7032a4154dc9056b76cf9625515e030e269cfbf8098fcfee7d/torch-2.12.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:e2ad3eb85d39c3cab62dfa93ed5a73516e6a53c6713cb97d004004fe089f0f1f", size = 532272265, upload-time = "2026-05-13T14:51:59.308Z" }, - { url = "https://files.pythonhosted.org/packages/f1/b4/92c80d1bbfee1c0036c06d1d2155a3065bd2423134c83bf8a47e65cd6b9b/torch-2.12.0-cp314-cp314-win_amd64.whl", hash = "sha256:c66696857e987efb8bc1777a37357ec4f60ab5e8af6250b83d6034437fa2d8f3", size = 122987138, upload-time = "2026-05-13T14:51:45.942Z" }, - { url = "https://files.pythonhosted.org/packages/7b/78/2e12b37ce50a19a037d7bc62d652a5a8f27385a7b05859d6bc9204f20cfe/torch-2.12.0-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:b4556715c8572758625d62b6e0ae3b1f76c440221913a6fb5e100f321fb4fb02", size = 88320100, upload-time = "2026-05-13T14:51:39.955Z" }, - { url = "https://files.pythonhosted.org/packages/56/5e/83c450ec7b0bb40a7b74611c1b5440f9260e33c54c90d556fd4a1f0fd955/torch-2.12.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:a43ac605a5e13116c72b64c359644cce0229f213dde48d2ae0ae5eb5becf7feb", size = 426391871, upload-time = "2026-05-13T14:52:14.989Z" }, - { url = "https://files.pythonhosted.org/packages/c9/e9/1a0b575d98d0afedd8f157d23fa3d2759421483660448e60d0a4b10b6daa/torch-2.12.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:6a7512adfdd7f6732e40de1c620831e3c75b39b98cef60b11d0c5f0a76473ec5", size = 532192241, upload-time = "2026-05-13T14:51:07.795Z" }, - { url = "https://files.pythonhosted.org/packages/88/21/afadd25ecd81b3cea1e11c73cf1ab41a983a50271548c3ec7ec3b9efc3e9/torch-2.12.0-cp314-cp314t-win_amd64.whl", hash = "sha256:5f96b63f8287f66a005dd1b5a6abba2920f11156c5e5c4d815f3e2050fd1aa16", size = 123231092, upload-time = "2026-05-13T14:51:18.854Z" }, + { url = "https://files.pythonhosted.org/packages/78/a9/97cbbc97002fff0de394a2da2cdfa859481fdca36996d7bd845d50aa9d8d/torch-2.6.0-cp311-cp311-manylinux1_x86_64.whl", hash = "sha256:7979834102cd5b7a43cc64e87f2f3b14bd0e1458f06e9f88ffa386d07c7446e1", size = 766715424, upload-time = "2025-01-29T16:25:15.874Z" }, + { url = "https://files.pythonhosted.org/packages/6d/fa/134ce8f8a7ea07f09588c9cc2cea0d69249efab977707cf67669431dcf5c/torch-2.6.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:ccbd0320411fe1a3b3fec7b4d3185aa7d0c52adac94480ab024b5c8f74a0bf1d", size = 95759416, upload-time = "2025-01-29T16:27:38.429Z" }, + { url = "https://files.pythonhosted.org/packages/11/c5/2370d96b31eb1841c3a0883a492c15278a6718ccad61bb6a649c80d1d9eb/torch-2.6.0-cp311-cp311-win_amd64.whl", hash = "sha256:46763dcb051180ce1ed23d1891d9b1598e07d051ce4c9d14307029809c4d64f7", size = 204164970, upload-time = "2025-01-29T16:26:16.182Z" }, + { url = "https://files.pythonhosted.org/packages/0b/fa/f33a4148c6fb46ca2a3f8de39c24d473822d5774d652b66ed9b1214da5f7/torch-2.6.0-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:94fc63b3b4bedd327af588696559f68c264440e2503cc9e6954019473d74ae21", size = 66530713, upload-time = "2025-01-29T16:26:38.881Z" }, + { url = "https://files.pythonhosted.org/packages/e5/35/0c52d708144c2deb595cd22819a609f78fdd699b95ff6f0ebcd456e3c7c1/torch-2.6.0-cp312-cp312-manylinux1_x86_64.whl", hash = "sha256:2bb8987f3bb1ef2675897034402373ddfc8f5ef0e156e2d8cfc47cacafdda4a9", size = 766624563, upload-time = "2025-01-29T16:23:19.084Z" }, + { url = "https://files.pythonhosted.org/packages/01/d6/455ab3fbb2c61c71c8842753b566012e1ed111e7a4c82e0e1c20d0c76b62/torch-2.6.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:b789069020c5588c70d5c2158ac0aa23fd24a028f34a8b4fcb8fcb4d7efcf5fb", size = 95607867, upload-time = "2025-01-29T16:25:55.649Z" }, + { url = "https://files.pythonhosted.org/packages/18/cf/ae99bd066571656185be0d88ee70abc58467b76f2f7c8bfeb48735a71fe6/torch-2.6.0-cp312-cp312-win_amd64.whl", hash = "sha256:7e1448426d0ba3620408218b50aa6ada88aeae34f7a239ba5431f6c8774b1239", size = 204120469, upload-time = "2025-01-29T16:24:01.821Z" }, + { url = "https://files.pythonhosted.org/packages/81/b4/605ae4173aa37fb5aa14605d100ff31f4f5d49f617928c9f486bb3aaec08/torch-2.6.0-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:9a610afe216a85a8b9bc9f8365ed561535c93e804c2a317ef7fabcc5deda0989", size = 66532538, upload-time = "2025-01-29T16:24:18.976Z" }, + { url = "https://files.pythonhosted.org/packages/24/85/ead1349fc30fe5a32cadd947c91bda4a62fbfd7f8c34ee61f6398d38fb48/torch-2.6.0-cp313-cp313-manylinux1_x86_64.whl", hash = "sha256:4874a73507a300a5d089ceaff616a569e7bb7c613c56f37f63ec3ffac65259cf", size = 766626191, upload-time = "2025-01-29T16:17:26.26Z" }, + { url = "https://files.pythonhosted.org/packages/dd/b0/26f06f9428b250d856f6d512413e9e800b78625f63801cbba13957432036/torch-2.6.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:a0d5e1b9874c1a6c25556840ab8920569a7a4137afa8a63a32cee0bc7d89bd4b", size = 95611439, upload-time = "2025-01-29T16:21:21.061Z" }, + { url = "https://files.pythonhosted.org/packages/c2/9c/fc5224e9770c83faed3a087112d73147cd7c7bfb7557dcf9ad87e1dda163/torch-2.6.0-cp313-cp313-win_amd64.whl", hash = "sha256:510c73251bee9ba02ae1cb6c9d4ee0907b3ce6020e62784e2d7598e0cfa4d6cc", size = 204126475, upload-time = "2025-01-29T16:21:55.394Z" }, + { url = "https://files.pythonhosted.org/packages/88/8b/d60c0491ab63634763be1537ad488694d316ddc4a20eaadd639cedc53971/torch-2.6.0-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:ff96f4038f8af9f7ec4231710ed4549da1bdebad95923953a25045dcf6fd87e2", size = 66536783, upload-time = "2025-01-29T16:22:08.559Z" }, ] [[package]] name = "torchvision" -version = "0.27.0" +version = "0.21.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy" }, @@ -4834,30 +4764,21 @@ dependencies = [ { name = "torch" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/cf/d6/a7e71e981042d5c573e2e61891b9023b190c88adb75b18bed8594371250c/torchvision-0.27.0-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:df0c166b6bdf7c47f88e81e8b43bc085451d5c50d0c5d1691bc474c1227d6fed", size = 1758812, upload-time = "2026-05-13T14:57:16.662Z" }, - { url = "https://files.pythonhosted.org/packages/93/f9/f542fb7e4476603fb237ebdc64369a7d11f18eb5a129aa2559cbdb710aee/torchvision-0.27.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:9bb9251f64b854124efed95d02953a89f7e2726c3ca662d7ea0151129157297f", size = 7831148, upload-time = "2026-05-13T14:57:08.37Z" }, - { url = "https://files.pythonhosted.org/packages/f6/61/7aa7cc2c9e8750027f6fb9ae3a7393ef43860bcdfe3966e2f71fee800e31/torchvision-0.27.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:f44453f107c296d5446a79f7ac59733ad8bf5ddfa04c53805dfbae298a42a798", size = 7575519, upload-time = "2026-05-13T14:56:50.552Z" }, - { url = "https://files.pythonhosted.org/packages/19/aa/929b358b1a643849b81ec95569938044cc37dc65ab10c84eb6d82fe1bfbb/torchvision-0.27.0-cp311-cp311-win_amd64.whl", hash = "sha256:b4aacff70ea4b7377f996f9048989c850d221fef33658ddbcae42aa5bd4ca11a", size = 3749475, upload-time = "2026-05-13T14:57:11.007Z" }, - { url = "https://files.pythonhosted.org/packages/9e/c8/5cd91932f7f3671b0743dc4ae1a4c16b1d0b45bf4087976277d325bda718/torchvision-0.27.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:1a6dd742a150645126df9e0b2e449874c1d635897c773b322c2e067e98382dfe", size = 1758824, upload-time = "2026-05-13T14:57:15.227Z" }, - { url = "https://files.pythonhosted.org/packages/d9/36/7fb7d19477b3d93283b52fea11fa8ee30ab9064a08c97b4a6b91445e26cb/torchvision-0.27.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:65772ff3ec4f4f5d680e30019835555dd239e7fefee4b0a846375fe1cb1592ef", size = 7831034, upload-time = "2026-05-13T14:57:06.483Z" }, - { url = "https://files.pythonhosted.org/packages/62/43/dfd894c3f8b01b5b33fde990f0159c1926ebc7b6e2c4193e2efb7da3c4cb/torchvision-0.27.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:7a9966a088d06b4cf6c610e03be62de469efa6f2cd2e7c7eed8e925ed6af59ac", size = 7579774, upload-time = "2026-05-13T14:56:59.337Z" }, - { url = "https://files.pythonhosted.org/packages/ff/0c/722e989f9cf026e97ef7cb24a9bb1859e099f72d247ae35388fb89729f73/torchvision-0.27.0-cp312-cp312-win_amd64.whl", hash = "sha256:2c037709072ca9b19750c0cbe9e8bb6f91c9a1be1befa26df33e281deccbd8c7", size = 4021073, upload-time = "2026-05-13T14:57:00.848Z" }, - { url = "https://files.pythonhosted.org/packages/d8/ae/36547812e6e047c1d80bcacd1b17a340612b08a6e876e0aabf3d0b9228b0/torchvision-0.27.0-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:41d6dae73e1af09fa82ded597ae57f2a2314285acde54b25890a8f8e51b999d7", size = 1758826, upload-time = "2026-05-13T14:57:05.262Z" }, - { url = "https://files.pythonhosted.org/packages/ae/30/32c4ea842738728a14e3df8c576c62dedcf5ae5cb6a5c984c6429ebe7524/torchvision-0.27.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:70f071c6f74b60d5fe8851636d8d4cd5f4fa29d57fd9348a87a6f17b990b95ba", size = 7789501, upload-time = "2026-05-13T14:56:57.786Z" }, - { url = "https://files.pythonhosted.org/packages/f6/24/4d0d48684251bd0673f87d633d5d88ab00227983b00591156eed2f86c8d5/torchvision-0.27.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:aaafa6962c9d91f42503de1957d6fa349907d028c06f335bd95da7a5bc57147d", size = 7579868, upload-time = "2026-05-13T14:56:41.618Z" }, - { url = "https://files.pythonhosted.org/packages/ba/da/e6edd051d2ba25adf23b120fa97f458dff888d098c51e84724f17d2d1470/torchvision-0.27.0-cp313-cp313-win_amd64.whl", hash = "sha256:aee384a2782c89517c4ab9061d2720ba59fd2ffe5ef89d0a149cc2d43abdf521", size = 4092700, upload-time = "2026-05-13T14:57:09.729Z" }, - { url = "https://files.pythonhosted.org/packages/fa/23/95dfa40431360f42ca949bf861434bed51164adfa8fb9801e05bf3194f50/torchvision-0.27.0-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:c5121f1b9ab09a7f73e837871deb8321551f7eaeb19d87aa00de9191968eae44", size = 1845008, upload-time = "2026-05-13T14:57:03.768Z" }, - { url = "https://files.pythonhosted.org/packages/23/b9/9dbdf76b2b49a75ba8088df6f7c755bdb520afb6c6dbac0102b46cde5e99/torchvision-0.27.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:1c01f0d1091ae22b9dfc082b0a0fe5faaf053686a29b4fb082ba7691375c73cf", size = 7791430, upload-time = "2026-05-13T14:56:56.206Z" }, - { url = "https://files.pythonhosted.org/packages/5c/6a/e4a16cf2f3310c2ea7760dc5d9054496844391e0f4c1fae87fefac2f3d9e/torchvision-0.27.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:dadea3c5ecfd05bbb2a3312ab0374f213c58bf6459cb059122e2f4dfe13d10ed", size = 7668441, upload-time = "2026-05-13T14:57:02.127Z" }, - { url = "https://files.pythonhosted.org/packages/00/70/01b6461117a6a94b5af3f8ee166bb0f045056f3cf187750c110dabfdfffa/torchvision-0.27.0-cp313-cp313t-win_amd64.whl", hash = "sha256:a49e55055a39a8506fe7e59850522cab004efb2c3839f6057658889c1d69c815", size = 4141602, upload-time = "2026-05-13T14:56:53.449Z" }, - { url = "https://files.pythonhosted.org/packages/92/22/c0633677b3b3f3e69554a21ac087bf705f829c40cd5e3783507b8c006681/torchvision-0.27.0-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:c1fac0fc2a7adf29481fc1938a0e7845c57ba1147a986784109c4d98f434ea8c", size = 1758818, upload-time = "2026-05-13T14:56:54.988Z" }, - { url = "https://files.pythonhosted.org/packages/48/e8/55f9d9667b56dae470e69e31beac9b00d458ea393feec1aae95cc4f3f1c9/torchvision-0.27.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:cbf89764fc76f3f17fbf80c12d5a89c691e91cb9d82c38412aaf0568655ffb19", size = 7789667, upload-time = "2026-05-13T14:56:48.858Z" }, - { url = "https://files.pythonhosted.org/packages/00/bc/6f8681daf3bbc4c315bb0005110f99d28e3ecd675bf9c8f2c0d393fbac7a/torchvision-0.27.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:91f61b9865423037c327eb56afa207cc72de874e458c361840db9dcf5ce0c0eb", size = 7579848, upload-time = "2026-05-13T14:56:38.209Z" }, - { url = "https://files.pythonhosted.org/packages/19/6c/8d8020e6bd1e46c53e487c9c4e9457a07f2ee28931028fb5d71e2da40adc/torchvision-0.27.0-cp314-cp314-win_amd64.whl", hash = "sha256:5bb82fc3c55daf1788621e504310b0a286f1069627a8742f692aebb075ef25a7", size = 4119284, upload-time = "2026-05-13T14:56:46.625Z" }, - { url = "https://files.pythonhosted.org/packages/8d/7e/e78c48662a8d551606efdbe11c6b9c1d6d2391b92cd0e4591b9e6a2412b8/torchvision-0.27.0-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:2c4099a15150143b9b034730b404a56d572efe0b79489b4c765d929cb4eac7f3", size = 1758828, upload-time = "2026-05-13T14:56:52.293Z" }, - { url = "https://files.pythonhosted.org/packages/21/dd/d03ee9f9ee7bf11a8c7c776fb8e7fd6102f59c013791a2a4e5175bd6cba7/torchvision-0.27.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:b4c6bb0a670dcba017b3643e21902c9b8a1cc1c127d602f1488fa29ec3c6e865", size = 7790618, upload-time = "2026-05-13T14:56:44.721Z" }, - { url = "https://files.pythonhosted.org/packages/39/08/4002336a74742be70728603ec1769feb2b55e0d19c532c9ec9f92008de76/torchvision-0.27.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:1c2db4bde82bc48ebff73436a6adf34d4f809448268a70d9a1285f5c8f92313d", size = 7580217, upload-time = "2026-05-13T14:56:43.274Z" }, - { url = "https://files.pythonhosted.org/packages/ed/cb/4dd4783eb3565f526ba6e64b6f6ca26c00eacc924cdfe60455db9d91b84b/torchvision-0.27.0-cp314-cp314t-win_amd64.whl", hash = "sha256:72bf547e58ddb948689734eed6f4b6a2031f979dba4fb08e3690688b392e929f", size = 4226392, upload-time = "2026-05-13T14:56:40.235Z" }, + { url = "https://files.pythonhosted.org/packages/4e/3d/b7241abfa3e6651c6e00796f5de2bd1ce4d500bf5159bcbfeea47e711b93/torchvision-0.21.0-1-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:ff96666b94a55e802ea6796cabe788541719e6f4905fc59c380fed3517b6a64d", size = 2329320, upload-time = "2025-03-18T17:25:52.272Z" }, + { url = "https://files.pythonhosted.org/packages/52/5b/76ca113a853b19c7b1da761f8a72cb6429b3bd0bf932537d8df4657f47c3/torchvision-0.21.0-1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:ffa2a16499508fe6798323e455f312c7c55f2a88901c9a7c0fb1efa86cf7e327", size = 2329878, upload-time = "2025-03-18T17:25:50.039Z" }, + { url = "https://files.pythonhosted.org/packages/4e/fe/5e193353706dab96fe73ae100d5a633ff635ce310e0d92f3bc2958d075b1/torchvision-0.21.0-1-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:7e9e9afa150e40cd2a8f0701c43cb82a8d724f512896455c0918b987f94b84a4", size = 2280711, upload-time = "2025-03-18T17:25:46.012Z" }, + { url = "https://files.pythonhosted.org/packages/29/88/00c69db213ee2443ada8886ec60789b227e06bb869d85ee324578221a7f7/torchvision-0.21.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:110d115333524d60e9e474d53c7d20f096dbd8a080232f88dddb90566f90064c", size = 1784141, upload-time = "2025-01-29T16:28:51.207Z" }, + { url = "https://files.pythonhosted.org/packages/be/a2/b0cedf0a411f1a5d75cfc0b87cde56dd1ddc1878be46a42c905cd8580220/torchvision-0.21.0-cp311-cp311-manylinux1_x86_64.whl", hash = "sha256:3891cd086c5071bda6b4ee9d266bb2ac39c998c045c2ebcd1e818b8316fb5d41", size = 7237719, upload-time = "2025-01-29T16:28:20.724Z" }, + { url = "https://files.pythonhosted.org/packages/8c/a1/ee962ef9d0b2bf7a6f8b14cb95acb70e05cd2101af521032a09e43f8582f/torchvision-0.21.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:54454923a50104c66a9ab6bd8b73a11c2fc218c964b1006d5d1fe5b442c3dcb6", size = 14700617, upload-time = "2025-01-29T16:28:30.247Z" }, + { url = "https://files.pythonhosted.org/packages/88/53/4ad334b9b1d8dd99836869fec139cb74a27781298360b91b9506c53f1d10/torchvision-0.21.0-cp311-cp311-win_amd64.whl", hash = "sha256:49bcfad8cfe2c27dee116c45d4f866d7974bcf14a5a9fbef893635deae322f2f", size = 1560523, upload-time = "2025-01-29T16:28:48.751Z" }, + { url = "https://files.pythonhosted.org/packages/6e/1b/28f527b22d5e8800184d0bc847f801ae92c7573a8c15979d92b7091c0751/torchvision-0.21.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:97a5814a93c793aaf0179cfc7f916024f4b63218929aee977b645633d074a49f", size = 1784140, upload-time = "2025-01-29T16:28:44.694Z" }, + { url = "https://files.pythonhosted.org/packages/36/63/0722e153fd27d64d5b0af45b5c8cb0e80b35a68cf0130303bc9a8bb095c7/torchvision-0.21.0-cp312-cp312-manylinux1_x86_64.whl", hash = "sha256:b578bcad8a4083b40d34f689b19ca9f7c63e511758d806510ea03c29ac568f7b", size = 7238673, upload-time = "2025-01-29T16:28:27.631Z" }, + { url = "https://files.pythonhosted.org/packages/bb/ea/03541ed901cdc30b934f897060d09bbf7a98466a08ad1680320f9ce0cbe0/torchvision-0.21.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:5083a5b1fec2351bf5ea9900a741d54086db75baec4b1d21e39451e00977f1b1", size = 14701186, upload-time = "2025-01-29T16:28:16.491Z" }, + { url = "https://files.pythonhosted.org/packages/4c/6a/c7752603060d076dfed95135b78b047dc71792630cbcb022e3693d6f32ef/torchvision-0.21.0-cp312-cp312-win_amd64.whl", hash = "sha256:6eb75d41e3bbfc2f7642d0abba9383cc9ae6c5a4ca8d6b00628c225e1eaa63b3", size = 1560520, upload-time = "2025-01-29T16:28:42.122Z" }, + { url = "https://files.pythonhosted.org/packages/f9/56/47d456b61c3bbce7bed4af3925c83d405bb87468e659fd3cf3d9840c3b51/torchvision-0.21.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:659b76c86757cb2ee4ca2db245e0740cfc3081fef46f0f1064d11adb4a8cee31", size = 1784141, upload-time = "2025-01-29T16:28:39.01Z" }, + { url = "https://files.pythonhosted.org/packages/cb/4c/99880813aa50e64447fb1c4c6c804a793d2d78f7f7c53e99ddee7fa175fa/torchvision-0.21.0-cp313-cp313-manylinux1_x86_64.whl", hash = "sha256:084ac3f5a1f50c70d630a488d19bf62f323018eae1b1c1232f2b7047d3a7b76d", size = 7238714, upload-time = "2025-01-29T16:28:25.658Z" }, + { url = "https://files.pythonhosted.org/packages/0b/2d/3c3ee10608310a395594aac7da8640372ed79c6585910ccae6919658dcdc/torchvision-0.21.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:5045a3a5f21ec3eea6962fa5f2fa2d4283f854caec25ada493fcf4aab2925467", size = 2281252, upload-time = "2025-01-29T16:28:40.687Z" }, + { url = "https://files.pythonhosted.org/packages/ed/b4/fc60e3bc003879d3de842baea258fffc3586f4b49cd435a5ba1e09c33315/torchvision-0.21.0-cp313-cp313-win_amd64.whl", hash = "sha256:9147f5e096a9270684e3befdee350f3cacafd48e0c54ab195f45790a9c146d67", size = 1560519, upload-time = "2025-01-29T16:28:22.527Z" }, ] [[package]] @@ -4891,7 +4812,7 @@ wheels = [ [[package]] name = "transformers" -version = "4.57.6" +version = "4.56.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "filelock" }, @@ -4905,28 +4826,19 @@ dependencies = [ { name = "tokenizers" }, { name = "tqdm" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/c4/35/67252acc1b929dc88b6602e8c4a982e64f31e733b804c14bc24b47da35e6/transformers-4.57.6.tar.gz", hash = "sha256:55e44126ece9dc0a291521b7e5492b572e6ef2766338a610b9ab5afbb70689d3", size = 10134912, upload-time = "2026-01-16T10:38:39.284Z" } +sdist = { url = "https://files.pythonhosted.org/packages/e9/40/2898c1a5f8b2fff2f94dffd856d6793320e590e1cfdc3a1453c31c364094/transformers-4.56.0.tar.gz", hash = "sha256:6ca9c3f38aa4da93ebf877db7156368c1c188c7465f09dbe70951e7622e987fa", size = 9840426, upload-time = "2025-08-29T18:23:46.247Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/03/b8/e484ef633af3887baeeb4b6ad12743363af7cce68ae51e938e00aaa0529d/transformers-4.57.6-py3-none-any.whl", hash = "sha256:4c9e9de11333ddfe5114bc872c9f370509198acf0b87a832a0ab9458e2bd0550", size = 11993498, upload-time = "2026-01-16T10:38:31.289Z" }, + { url = "https://files.pythonhosted.org/packages/98/41/6ba7c2eafa839069f1f949c30cd22e791b14b90c7af8cb0c65cc47702dca/transformers-4.56.0-py3-none-any.whl", hash = "sha256:bacf539c38dd850690856881c4974321af93a22f2ee96bcc994741a2121d8e71", size = 11607938, upload-time = "2025-08-29T18:23:33.92Z" }, ] [[package]] name = "triton" -version = "3.7.0" +version = "3.2.0" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/b8/c1/5d842314bb6c78442cc60437928781701c6050b8d479bc2a1aed691d37ca/triton-3.7.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a9e71fc392675fac364e0ecf4ef3f76f85b7f5433a16f4c3c5fe5f05a52c85fe", size = 188480277, upload-time = "2026-05-07T19:05:03.231Z" }, - { url = "https://files.pythonhosted.org/packages/13/31/8315ea5f8dd18e60970b3022e3a8b93fd37e0b784fbbef86e10c8e6e5ca1/triton-3.7.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:22bacffce443f54593dd20f05294d5a40622e0ea9ab632816f87154504356221", size = 201415942, upload-time = "2026-05-07T18:46:06.479Z" }, - { url = "https://files.pythonhosted.org/packages/f7/13/ec05adfcd87311d532ba61e3af143e8be59fcd26675884c4682841406a20/triton-3.7.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a4bf49b00a7a377a68a6da603a876e797614e6455a80e9021669c476a953ad9a", size = 188505104, upload-time = "2026-05-07T19:05:09.843Z" }, - { url = "https://files.pythonhosted.org/packages/62/7b/468a576e35beef1426e0828e28e9ba9e65f5474d496f16ee126c15646324/triton-3.7.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8f111161d49bf903c0eaedde3962353a3d841c08a836839b7cc1025b8426efcf", size = 201457567, upload-time = "2026-05-07T18:46:13.505Z" }, - { url = "https://files.pythonhosted.org/packages/01/e1/a59a583de59b8f62c495d67c80ee3ea97d09e91ac80c4c6e76456ed8d8ac/triton-3.7.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:abdf6beaa89b1bcfb9a43cd990536ce66091a997841a4814b260b7bee4c88c3c", size = 188503209, upload-time = "2026-05-07T19:05:17.935Z" }, - { url = "https://files.pythonhosted.org/packages/30/b1/b7507bb9815d403927c8dd51d4158ed2e11751a92dbc118a044f247b6848/triton-3.7.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a35d7afe3f3f058e7ec49fcce09794049e0ffc5c59019ac25ec3413741b8c4e7", size = 201453566, upload-time = "2026-05-07T18:46:20.427Z" }, - { url = "https://files.pythonhosted.org/packages/a6/8f/0bea7a6a0c989315c9135a1d7fb37e41905cfb3a17cbc1f10044ebd4cc3a/triton-3.7.0-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cc1d61c172d257db80ddf42595131fb196ad2e9bdd751e90fe2ef13531734e8b", size = 188612899, upload-time = "2026-05-07T19:05:24.955Z" }, - { url = "https://files.pythonhosted.org/packages/e1/02/d96f57828d0912aec733b9bc7e0e7dbfd2c6f079a8fa433ac25cb93d1a30/triton-3.7.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:70fb9bbdc9f400afc54bbf6eb2670af28829a6ae3996863317964783141daf56", size = 201553816, upload-time = "2026-05-07T18:46:27.49Z" }, - { url = "https://files.pythonhosted.org/packages/40/fb/82a802dac4689f2a2fb2e69302e6a138eecc3e175bbe976ba3cfc717683a/triton-3.7.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c4a44a8476d0d3571eac4e4d1048e1ff75aad81a09ff4602ccfc56c6dea1672e", size = 188507879, upload-time = "2026-05-07T19:05:32.209Z" }, - { url = "https://files.pythonhosted.org/packages/8f/af/9904ec6d3c93d9b24e5ec360445bbdf758b7f00bfbeedb89cb0eb64eb8bb/triton-3.7.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b9b85e72968a9d8bba5ddb24e9b64aaabaf48affb042f2755cb7cfa92b7531ce", size = 201460637, upload-time = "2026-05-07T18:46:34.749Z" }, - { url = "https://files.pythonhosted.org/packages/a1/f9/4835a8ea746b88727d8899f4e3ccce4f9cacb38abfc3bb0a638266c53111/triton-3.7.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:18a160de426fd99f92b0baf509045360afbd3bfaa0b4a5171dde800ec9f09684", size = 188608706, upload-time = "2026-05-07T19:05:39.218Z" }, - { url = "https://files.pythonhosted.org/packages/c1/68/fa86e5a39608000f645535b2c124920126327ab731f8c4fafd5b07ff8d4b/triton-3.7.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ce061073102714b725f3660ec6939d94a1da7984b3aa99c921417cae273672f5", size = 201546766, upload-time = "2026-05-07T18:46:42.088Z" }, + { url = "https://files.pythonhosted.org/packages/a7/2e/757d2280d4fefe7d33af7615124e7e298ae7b8e3bc4446cdb8e88b0f9bab/triton-3.2.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8009a1fb093ee8546495e96731336a33fb8856a38e45bb4ab6affd6dbc3ba220", size = 253157636, upload-time = "2025-01-22T19:12:51.322Z" }, + { url = "https://files.pythonhosted.org/packages/06/00/59500052cb1cf8cf5316be93598946bc451f14072c6ff256904428eaf03c/triton-3.2.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8d9b215efc1c26fa7eefb9a157915c92d52e000d2bf83e5f69704047e63f125c", size = 253159365, upload-time = "2025-01-22T19:13:24.648Z" }, + { url = "https://files.pythonhosted.org/packages/c7/30/37a3384d1e2e9320331baca41e835e90a3767303642c7a80d4510152cbcf/triton-3.2.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e5dfa23ba84541d7c0a531dfce76d8bcd19159d50a4a8b14ad01e91734a5c1b0", size = 253154278, upload-time = "2025-01-22T19:13:54.221Z" }, ] [[package]] @@ -5206,6 +5118,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/38/5f/722007afc8998343720b5ccf2a619c6d609ac029bac5390c532ce28433fa/xarray_spatial-0.10.6-py3-none-any.whl", hash = "sha256:c6ae298424e016cea09a9477cd6fcb518231bc839d70e505bc0f0f104fed44d8", size = 6131730, upload-time = "2026-06-09T01:13:54.171Z" }, ] +[[package]] +name = "xformers" +version = "0.0.29.post3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "torch", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c1/fd/e9201fbee6a1a6d7a9c67c24a256ad4c2377bc67a634f7dbeaea23bd668a/xformers-0.0.29.post3.tar.gz", hash = "sha256:0b77c67ecc3c9fdd8a0e4399e675adf12e2ff40285e00974cca2d09108157f60", size = 8461348, upload-time = "2025-02-10T18:33:04.046Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/55/05/9c9faf1c7b3b7b986bbf7a488a185eb67670a8435d0eae94aa59f56181cd/xformers-0.0.29.post3-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:bbf2f500dfdbcf4649bf568cc2c9f434399f704dc4064fd1fbdbef2b524a8139", size = 43362399, upload-time = "2025-02-10T18:31:38.988Z" }, + { url = "https://files.pythonhosted.org/packages/2d/4a/20b2d9ac50efa0d40fbdb13283fd168cc2db28a2f21a159abbdd17a24213/xformers-0.0.29.post3-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:08fa92f3e06372c4ce2a5306c54ae3d4a3a399fc7e24e02aac3761112ec3aeed", size = 43364118, upload-time = "2025-02-10T18:32:10.063Z" }, +] + [[package]] name = "yarl" version = "1.24.2" From 4a70e5aad680eafa4bee7eb2d70f6910721f1add Mon Sep 17 00:00:00 2001 From: Craig Myles Date: Mon, 3 Aug 2026 14:18:16 +0100 Subject: [PATCH 2/4] Address KRONOS2 review findings --- CHANGELOG.md | 7 ++ CITATION.cff | 1 - README.md | 14 +-- docs/MODALITIES.md | 5 + docs/MODELS.md | 10 +- docs/usage.md | 10 +- pyproject.toml | 10 +- .../embedders/kronos2_embedder.py | 91 +++++++++++++++++-- src/raw2features/embedders/registry.yaml | 11 ++- src/raw2features/pipeline/runner.py | 3 +- src/raw2features/segmenters/nuclear.py | 24 +++-- tests/test_config_integrity.py | 32 +++++++ tests/test_kronos2.py | 76 +++++++++++++++- tests/test_multiplex.py | 61 +++++++++++++ tests/test_release_metadata.py | 43 ++++++++- uv.lock | 8 +- 16 files changed, 358 insertions(+), 48 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 745d221..aa84ac2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,6 +19,13 @@ layout, geometry-only `grid_hash`, and brightfield execution path are unchanged. 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. - 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. diff --git a/CITATION.cff b/CITATION.cff index 12c2fa4..d123f87 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -21,4 +21,3 @@ keywords: - feature extraction license: MIT version: "0.2.1" -date-released: 2026-07-29 diff --git a/README.md b/README.md index 5d799c6..c7229c6 100644 --- a/README.md +++ b/README.md @@ -92,12 +92,14 @@ 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, the extra installs -the authors' pinned Torch/xFormers/Transformers/timm reference stack; raw2features -validates this path on Python 3.12. 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 +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/)): diff --git a/docs/MODALITIES.md b/docs/MODALITIES.md index 7e34a84..ef9c1d8 100644 --- a/docs/MODALITIES.md +++ b/docs/MODALITIES.md @@ -88,6 +88,11 @@ raw2features embed SLIDE.ome.zarr OUT -m kronos2 \ 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 diff --git a/docs/MODELS.md b/docs/MODELS.md index ec2a768..ae8f519 100644 --- a/docs/MODELS.md +++ b/docs/MODELS.md @@ -271,10 +271,12 @@ back to **1.0 µm/px** unless run alongside a model that supplies one or given a 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. The `[kronos2]` extra - installs the authors' pinned Torch, xFormers, Transformers, and timm reference stack - on x86-64 Linux and is validated with Python 3.12. The pinned Torch release does not - support Python 3.14. `--compile` is not supported for KRONOS2 in v0.2.1. + 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 diff --git a/docs/usage.md b/docs/usage.md index bbe87b5..cb26c41 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -358,10 +358,12 @@ Upstream recommends batches of at least eight for reproducibility. raw2features fixed eight-sample forward microbatches and pads a shorter final microbatch before discarding the padded outputs. This runtime contract is recorded. The global `--batch-size` may be left unchanged; it does not override this internal cap. Bitwise -equality across different hardware is not promised. The `[kronos2]` extra installs the -authors' pinned x86-64 Linux reference stack and is validated with Python 3.12. The -pinned Torch release does not support Python 3.14. `--compile` is not supported for -KRONOS2 in v0.2.1; use the validated eager path. +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 xFormers reference path; +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; use the validated eager path. The released KRONOS2 repository includes 288 usable marker-metadata entries, 268 marked as pretraining markers, and stores their text vectors in the checkpoint. Normal diff --git a/pyproject.toml b/pyproject.toml index b15e335..1c7eff1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -143,13 +143,13 @@ kronos = ["raw2features[torch,models]", "tifffile>=2024", "imagecodecs"] # KRONOS2 is distributed as a gated Transformers custom-code snapshot. Its authors # publish this exact stack for reproducible inference, so the model-specific extra # deliberately pins it rather than silently substituting newer attention kernels. -# xFormers is part of that reference path on Linux; the upstream code has a slower, -# more memory-intensive fallback elsewhere. BioLinkBERT is not loaded for published -# markers. The reference stack is validated on Python 3.12; its pinned Torch release -# does not provide Python 3.14 wheels. +# xFormers is part of that reference path where its pinned release has a wheel; the +# upstream code has a slower, more memory-intensive PyTorch fallback elsewhere. +# BioLinkBERT is not loaded for published markers. The reference stack is validated on +# Python 3.12; its pinned Torch release does not provide Python 3.14 wheels. kronos2 = [ "torch==2.6.0", - "xformers==0.0.29.post3; sys_platform == 'linux' and platform_machine == 'x86_64'", + "xformers==0.0.29.post3; sys_platform == 'linux' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and python_version < '3.13'", "transformers==4.56.0", "timm==1.0.19", "huggingface-hub>=0.28.1,<1.0", diff --git a/src/raw2features/embedders/kronos2_embedder.py b/src/raw2features/embedders/kronos2_embedder.py index 300b022..5752865 100644 --- a/src/raw2features/embedders/kronos2_embedder.py +++ b/src/raw2features/embedders/kronos2_embedder.py @@ -2,10 +2,11 @@ KRONOS2 is distributed as a gated Hugging Face repository containing custom ``transformers`` code, a marker-aware DINOv2 ViT-B/16 checkpoint, and the marker -metadata that defines both normalisation and marker identity. raw2features downloads -one immutable snapshot, verifies both content-affecting data artifacts, and gives only -that local directory to ``AutoModel``. This prevents the pinned custom loader from -silently resolving a newer copy of its own weights. +metadata that defines both normalisation and marker identity. raw2features downloads +only the files required for inference from one immutable revision into a managed cache, +verifies both content-affecting data artifacts, and gives only that local directory to +``AutoModel``. This prevents the pinned custom loader from silently resolving a newer +copy of its own weights. Published markers use text embeddings already baked into the checkpoint. The upstream ``register_additional_markers`` path runs BioLinkBERT for genuinely novel markers. @@ -25,9 +26,11 @@ import json import math import os +import platform import sys import tempfile import unicodedata +from importlib import metadata as importlib_metadata from pathlib import Path from typing import TYPE_CHECKING, Any @@ -56,6 +59,7 @@ KRONOS2_MARKER_METADATA_FILENAME = "marker_metadata.csv" +KRONOS2_XFORMERS_VERSION = "0.0.29.post3" KRONOS2_SNAPSHOT_TOP_LEVEL_ALLOWLIST = ( "config.json", "configuration_kronos2.py", @@ -103,6 +107,22 @@ "output": "slice_back_to_unpadded_row_count", }, "output_order": "input_order", + "attention_backend_policy": { + "selection": "pinned_xformers_in_reference_environment_else_forced_fallback", + "pinned_reference": { + "implementation": "xformers.ops.memory_efficient_attention", + "package_version": KRONOS2_XFORMERS_VERSION, + "install_environment": "linux_x86_64_cpython_3_11_or_3_12", + }, + "fallback": { + "implementation": "upstream_dinov2_pytorch_attention", + "condition": ( + "outside_reference_environment_or_xformers_unavailable_or_disabled" + ), + "enforcement": "temporary_XFORMERS_DISABLED_during_construction", + "preimport_conflict": "fail_closed", + }, + }, } KRONOS2_MARKER_MATCHING_CONTRACT: dict[str, Any] = { "name": "upstream_separator_insensitive_exact", @@ -303,11 +323,50 @@ def _assert_allowlisted_runtime_snapshot(snapshot: str) -> None: unexpected = sorted(actual - allowed) if unexpected: raise ValueError( - "KRONOS2's isolated runtime snapshot contains unexpected top-level " + "KRONOS2's raw2features-managed cache contains unexpected top-level " f"entries: {unexpected}. Refusing to add it to the Python import path." ) +def _allows_pinned_xformers( + *, + system: str | None = None, + machine: str | None = None, + implementation: str | None = None, + python_version: tuple[int, int] | None = None, +) -> bool: + """Whether this process may use the one validated xFormers distribution.""" + + runtime = tuple(sys.version_info[:2]) if python_version is None else python_version + if ( + (system or platform.system()).casefold() != "linux" + or (machine or platform.machine()).casefold() not in {"x86_64", "amd64"} + or (implementation or platform.python_implementation()).casefold() != "cpython" + or runtime not in {(3, 11), (3, 12)} + ): + return False + try: + return importlib_metadata.version("xformers") == KRONOS2_XFORMERS_VERSION + except importlib_metadata.PackageNotFoundError: + return False + + +def _assert_forced_fallback_is_not_preimported_with_xformers() -> None: + """Fail if an already-imported upstream module makes fallback unenforceable.""" + + conflicts = [ + name + for name in ("dinov2.layers.attention", "dinov2.layers.block") + if bool(getattr(sys.modules.get(name), "XFORMERS_AVAILABLE", False)) + ] + if conflicts: + raise RuntimeError( + "KRONOS2 must use its PyTorch attention fallback in this environment, " + "but an already-imported DINOv2 module enabled xFormers: " + f"{conflicts}. Start a fresh Python process before loading KRONOS2." + ) + + @register("embedders", "kronos2") class Kronos2Embedder(Embedder): """KRONOS2 marker-aware ViT-B/16, returning the 768-dimensional CLS token.""" @@ -364,8 +423,14 @@ def load( # snapshot_download call. The outer Transformers resolver also receives # local_files_only=True. Upstream temporarily inserts the directory at # sys.path[0] for its bundled dinov2 package; remove that new entry after - # construction so its module names cannot shadow later imports. - snapshot_was_on_sys_path = snapshot in sys.path + # construction so the directory is not left at the front of the caller's + # import search path. + force_attention_fallback = not _allows_pinned_xformers() + previous_xformers_disabled = os.environ.get("XFORMERS_DISABLED") + if force_attention_fallback: + _assert_forced_fallback_is_not_preimported_with_xformers() + os.environ["XFORMERS_DISABLED"] = "1" + snapshot_path_count = sys.path.count(snapshot) try: model = AutoModel.from_pretrained( snapshot, @@ -374,8 +439,16 @@ def load( device=device, ) finally: - if not snapshot_was_on_sys_path: - sys.path[:] = [entry for entry in sys.path if entry != snapshot] + added_count = max(0, sys.path.count(snapshot) - snapshot_path_count) + for _ in range(added_count): + # Upstream inserts at the front, and list.remove() removes the first + # occurrence. Preserve any entry that belonged to the caller. + sys.path.remove(snapshot) + if force_attention_fallback: + if previous_xformers_disabled is None: + os.environ.pop("XFORMERS_DISABLED", None) + else: + os.environ["XFORMERS_DISABLED"] = previous_xformers_disabled model = model.float().eval().to(device) self._model = model self._device = device diff --git a/src/raw2features/embedders/registry.yaml b/src/raw2features/embedders/registry.yaml index 08c2b20..4b71142 100644 --- a/src/raw2features/embedders/registry.yaml +++ b/src/raw2features/embedders/registry.yaml @@ -1037,11 +1037,12 @@ kronos2: marker_metadata_sha256: "1e6791448a792c0edec9b2ba74efceaefba84d96852351cdc58d67f80b539074" notes: >- KRONOS2 (MahmoodLab) is a native marker-aware multiplex ViT-B/16 that returns - a 768-d CLS vector. The complete pinned Hugging Face snapshot is loaded locally - after the weight and marker-metadata hashes are verified. Marker names match - exactly after the upstream separator-only normalization; no biological aliases - or KRONOSv1 compound-token fallback is applied. Published marker text embeddings - are checkpoint buffers, so ordinary inference does not download BioLinkBERT. + a 768-d CLS vector. Only the inference files allowlisted from the pinned Hugging + Face revision are placed in a raw2features-managed cache. The weight and marker + metadata hashes are verified before its custom code executes. Marker names match + exactly after the upstream separator-only normalization; no biological aliases or + KRONOSv1 compound-token fallback is applied. Published marker text embeddings are + checkpoint buffers, so ordinary inference does not download BioLinkBERT. An explicit complete additional-marker CSV enables the upstream registration path with user-supplied statistics and a separately pinned, SHA-verified BioLinkBERT snapshot; that conditional contract is fingerprinted. Requires the `kronos2` diff --git a/src/raw2features/pipeline/runner.py b/src/raw2features/pipeline/runner.py index 45b94e5..fc36e11 100644 --- a/src/raw2features/pipeline/runner.py +++ b/src/raw2features/pipeline/runner.py @@ -2694,7 +2694,8 @@ def _segment(reader, cfg: RunConfig, segmenter_name: str | None = None): """Return (tissue, seg_meta); (None, none-meta) when segmentation is off. ``segmenter_name`` overrides ``cfg.segmenter`` (multiplex slides route to the - ``nuclear`` segmenter, which thresholds the DAPI/Hoechst channel). + ``nuclear`` segmenter, which thresholds an established DAPI/Hoechst/DNA + channel or a DRAQ5 fallback). """ if cfg.no_seg: return None, {"segmenter": "none"} diff --git a/src/raw2features/segmenters/nuclear.py b/src/raw2features/segmenters/nuclear.py index 19dd15d..060fce5 100644 --- a/src/raw2features/segmenters/nuclear.py +++ b/src/raw2features/segmenters/nuclear.py @@ -1,9 +1,9 @@ """Nuclear-channel tissue segmenter for named-channel multiplex slides. Multiplex slides have no H&E saturation/colour to threshold; instead, tissue is where -the nuclear stain (DAPI / Hoechst) is present. This segmenter finds the nuclear marker -channel(s) by name (via the reader's ``channel_names``), reads them at a low-res level, -and runs Otsu + morphology - the multiplex analogue of the default +a nuclear stain is present. This segmenter finds established DAPI/Hoechst/DNA marker +channels by name, falling back to DRAQ5 when none is present, reads them at a low-res +level, and runs Otsu + morphology - the multiplex analogue of the default Otsu-on-saturation segmenter. Pure OpenCV + numpy. """ @@ -24,16 +24,20 @@ if TYPE_CHECKING: from raw2features.readers.base import WSISource -_NUCLEAR_ALIASES = ("dapi", "hoechst", "hochst", "dna") +_NUCLEAR_ALIASES = ("dapi", "hoechst", "hochst", "dna", "draq5") # Numbered DNA channels occur both as plain labels (``DNA1``) and as delimited # tokens in IMC exports (``191Ir_DNA1`` / ``Ir193_DNA2``). Keep the boundaries # strict so biomarkers and prose containing ``DNA`` never match. _DNA_NUMBERED_RE = re.compile(r"(? None: assert raw2features.__version__ == version assert citation["version"] == version - assert f"## [{version}]" in (ROOT / "CHANGELOG.md").read_text() + changelog = (ROOT / "CHANGELOG.md").read_text() + assert f"## [{version}]" in changelog + if f"## [{version}] - Unreleased" in changelog: + assert "date-released" not in citation + + +def test_kronos2_xformers_pin_only_targets_published_wheel_versions() -> None: + project = tomllib.loads((ROOT / "pyproject.toml").read_text()) + requirement = Requirement( + next( + item + for item in project["project"]["optional-dependencies"]["kronos2"] + if item.startswith("xformers") + ) + ) + assert str(requirement.specifier) == "==0.0.29.post3" + assert requirement.marker is not None + + def selected(python_version: str) -> bool: + return requirement.marker.evaluate( + { + "sys_platform": "linux", + "platform_machine": "x86_64", + "platform_python_implementation": "CPython", + "python_version": python_version, + "python_full_version": f"{python_version}.0", + } + ) + + assert selected("3.11") + assert selected("3.12") + assert not selected("3.13") + + pypy_environment = { + "sys_platform": "linux", + "platform_machine": "x86_64", + "platform_python_implementation": "PyPy", + "python_version": "3.12", + "python_full_version": "3.12.0", + } + assert not requirement.marker.evaluate(pypy_environment) diff --git a/uv.lock b/uv.lock index 3380dac..af8a22a 100644 --- a/uv.lock +++ b/uv.lock @@ -3396,7 +3396,7 @@ kronos2 = [ { name = "timm" }, { name = "torch" }, { name = "transformers" }, - { name = "xformers", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "xformers", marker = "python_full_version < '3.13' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux'" }, ] madeleine = [ { name = "einops" }, @@ -3584,7 +3584,7 @@ requires-dist = [ { name = "transformers", marker = "extra == 'kronos2'", specifier = "==4.56.0" }, { name = "transformers", marker = "extra == 'models'", specifier = ">=4.40" }, { name = "typer", specifier = ">=0.12" }, - { name = "xformers", marker = "platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'kronos2'", specifier = "==0.0.29.post3" }, + { name = "xformers", marker = "python_full_version < '3.13' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux' and extra == 'kronos2'", specifier = "==0.0.29.post3" }, { name = "zarr", marker = "extra == 'zarr'", specifier = ">=3.1" }, ] provides-extras = ["all", "zarr", "image", "stain", "torch", "models", "conch", "hibou", "kronos", "kronos2", "musk", "open-clip", "grandqc", "prism", "madeleine", "gigapath-slide", "chief", "tangle", "seal", "spatialdata", "h5", "tables", "s3", "benchmark"] @@ -5123,8 +5123,8 @@ name = "xformers" version = "0.0.29.post3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "numpy", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "torch", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "numpy", marker = "python_full_version < '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "torch", marker = "python_full_version < '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/c1/fd/e9201fbee6a1a6d7a9c67c24a256ad4c2377bc67a634f7dbeaea23bd668a/xformers-0.0.29.post3.tar.gz", hash = "sha256:0b77c67ecc3c9fdd8a0e4399e675adf12e2ff40285e00974cca2d09108157f60", size = 8461348, upload-time = "2025-02-10T18:33:04.046Z" } wheels = [ From 94d488f1c83adc691b19df04eaf45f5e3055b813 Mon Sep 17 00:00:00 2001 From: Craig Myles Date: Mon, 3 Aug 2026 18:52:32 +0100 Subject: [PATCH 3/4] Fingerprint KRONOS2 attention backend --- CHANGELOG.md | 2 + src/raw2features/embedders/fingerprint.py | 18 ++++- .../embedders/kronos2_embedder.py | 34 ++++++++- src/raw2features/pipeline/runner.py | 15 +++- tests/test_kronos2.py | 75 ++++++++++++++++++- 5 files changed, 139 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index aa84ac2..9a6bb4a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,6 +26,8 @@ layout, geometry-only `grid_hash`, and brightfield execution path are unchanged. - 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 + produced by the two paths cannot be silently combined during resume. - 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. diff --git a/src/raw2features/embedders/fingerprint.py b/src/raw2features/embedders/fingerprint.py index e082a92..fa4b5a0 100644 --- a/src/raw2features/embedders/fingerprint.py +++ b/src/raw2features/embedders/fingerprint.py @@ -472,7 +472,12 @@ def _patch_constructor(spec: ModelSpec) -> dict[str, Any]: return constructor -def patch_output_fingerprint(spec: ModelSpec, resolved_amp: str) -> dict[str, Any]: +def patch_output_fingerprint( + spec: ModelSpec, + resolved_amp: str, + *, + resolved_attention_backend: Mapping[str, Any] | None = None, +) -> dict[str, Any]: """Fingerprint the complete persisted-output contract for one patch model.""" if spec.multiplex is not None: @@ -567,11 +572,20 @@ def patch_output_fingerprint(spec: ModelSpec, resolved_amp: str) -> dict[str, An } ) if spec.family == "kronos2": - from .kronos2_embedder import KRONOS2_SCALING_CONTRACT + from .kronos2_embedder import ( + KRONOS2_SCALING_CONTRACT, + resolved_kronos2_attention_backend, + ) payload["preprocessing"]["native_multiplex"] = deepcopy( KRONOS2_SCALING_CONTRACT ) + backend = ( + resolved_kronos2_attention_backend() + if resolved_attention_backend is None + else dict(resolved_attention_backend) + ) + payload["output"]["attention_backend"] = deepcopy(backend) return make_output_fingerprint(payload) diff --git a/src/raw2features/embedders/kronos2_embedder.py b/src/raw2features/embedders/kronos2_embedder.py index 5752865..733f163 100644 --- a/src/raw2features/embedders/kronos2_embedder.py +++ b/src/raw2features/embedders/kronos2_embedder.py @@ -60,6 +60,15 @@ KRONOS2_MARKER_METADATA_FILENAME = "marker_metadata.csv" KRONOS2_XFORMERS_VERSION = "0.0.29.post3" +KRONOS2_XFORMERS_ATTENTION_BACKEND: dict[str, Any] = { + "implementation": "xformers.ops.memory_efficient_attention", + "package": "xformers", + "package_version": KRONOS2_XFORMERS_VERSION, +} +KRONOS2_PYTORCH_ATTENTION_BACKEND: dict[str, Any] = { + "implementation": "upstream_dinov2_pytorch_attention", + "package": "torch", +} KRONOS2_SNAPSHOT_TOP_LEVEL_ALLOWLIST = ( "config.json", "configuration_kronos2.py", @@ -351,6 +360,22 @@ def _allows_pinned_xformers( return False +def resolved_kronos2_attention_backend() -> dict[str, Any]: + """Return the attention implementation this process will use for KRONOS2.""" + + if "XFORMERS_DISABLED" in os.environ or not _allows_pinned_xformers(): + return dict(KRONOS2_PYTORCH_ATTENTION_BACKEND) + imported_flags = [ + bool(module.XFORMERS_AVAILABLE) + for name in ("dinov2.layers.attention", "dinov2.layers.block") + if (module := sys.modules.get(name)) is not None + and hasattr(module, "XFORMERS_AVAILABLE") + ] + if imported_flags and not all(imported_flags): + return dict(KRONOS2_PYTORCH_ATTENTION_BACKEND) + return dict(KRONOS2_XFORMERS_ATTENTION_BACKEND) + + def _assert_forced_fallback_is_not_preimported_with_xformers() -> None: """Fail if an already-imported upstream module makes fallback unenforceable.""" @@ -425,7 +450,10 @@ def load( # sys.path[0] for its bundled dinov2 package; remove that new entry after # construction so the directory is not left at the front of the caller's # import search path. - force_attention_fallback = not _allows_pinned_xformers() + attention_backend = resolved_kronos2_attention_backend() + force_attention_fallback = ( + attention_backend != KRONOS2_XFORMERS_ATTENTION_BACKEND + ) previous_xformers_disabled = os.environ.get("XFORMERS_DISABLED") if force_attention_fallback: _assert_forced_fallback_is_not_preimported_with_xformers() @@ -453,6 +481,10 @@ def load( self._model = model self._device = device self._dtype = torch.float32 + # Resolve once more after upstream imports DINOv2. A broken xFormers binary + # can make upstream fall back even when the pinned distribution is installed; + # the loaded-model contract must record what its module actually selected. + self._resolved_attention_backend = resolved_kronos2_attention_backend() self._metadata_path = metadata self._load_marker_metadata(metadata) self._panel: dict[str, Any] | None = None diff --git a/src/raw2features/pipeline/runner.py b/src/raw2features/pipeline/runner.py index fc36e11..c360d66 100644 --- a/src/raw2features/pipeline/runner.py +++ b/src/raw2features/pipeline/runner.py @@ -2530,9 +2530,22 @@ def _loaded_model_contracts( selected, getattr(embedder, "_device", "cpu"), ) + fingerprint_kwargs = {} + if embedder.spec.family == "kronos2": + backend = getattr(embedder, "_resolved_attention_backend", None) + if backend is None: + raise RuntimeError( + "Loaded KRONOS2 embedder has no resolved attention-backend " + "contract. Reload the model before writing outputs." + ) + fingerprint_kwargs["resolved_attention_backend"] = backend contracts[embedder.name] = { "embedding_dim": int(embedder.embedding_dim), - "output_fingerprint": patch_output_fingerprint(embedder.spec, effective), + "output_fingerprint": patch_output_fingerprint( + embedder.spec, + effective, + **fingerprint_kwargs, + ), } return _bind_native_multiplex_panel_contracts( contracts, diff --git a/tests/test_kronos2.py b/tests/test_kronos2.py index fea15b6..bfc9cea 100644 --- a/tests/test_kronos2.py +++ b/tests/test_kronos2.py @@ -23,21 +23,28 @@ KRONOS2_BIOLINKBERT_SOURCE, KRONOS2_FORWARD_CONTRACT, KRONOS2_MARKER_MATCHING_CONTRACT, + KRONOS2_PYTORCH_ATTENTION_BACKEND, KRONOS2_SCALING_CONTRACT, KRONOS2_SNAPSHOT_ALLOW_PATTERNS, KRONOS2_SNAPSHOT_TOP_LEVEL_ALLOWLIST, + KRONOS2_XFORMERS_ATTENTION_BACKEND, KRONOS2_XFORMERS_VERSION, Kronos2Embedder, _allows_pinned_xformers, _marker_match_key, _scaling_factor, + resolved_kronos2_attention_backend, ) from raw2features.embedders.kronos2_metadata import ( KRONOS2_ADDITIONAL_MARKER_COLUMNS, parse_kronos2_additional_markers, ) from raw2features.embedders.model_registry import get_spec -from raw2features.pipeline.runner import RunConfig, expected_model_contracts +from raw2features.pipeline.runner import ( + RunConfig, + _loaded_model_contracts, + expected_model_contracts, +) def _digest(data: bytes) -> str: @@ -190,6 +197,7 @@ def test_kronos2_registry_and_fingerprint_cover_the_complete_static_contract(): "reg_tokens": 0, "modality": "multiplex", "resolved_amp": "fp32", + "attention_backend": resolved_kronos2_attention_backend(), } assert KRONOS2_BIOLINKBERT_SOURCE not in json.dumps(payload, sort_keys=True) with pytest.raises(ValueError, match="does not accept requested AMP"): @@ -741,6 +749,70 @@ def test_only_the_pinned_reference_xformers_environment_is_allowed( ) +def test_resolved_attention_backend_changes_the_output_fingerprint(monkeypatch): + import raw2features.embedders.kronos2_embedder as module + + spec = get_spec("kronos2") + monkeypatch.delenv("XFORMERS_DISABLED", raising=False) + monkeypatch.setattr(module, "_allows_pinned_xformers", lambda: True) + + xformers = patch_output_fingerprint(spec, "fp32") + assert resolved_kronos2_attention_backend() == (KRONOS2_XFORMERS_ATTENTION_BACKEND) + assert xformers["payload"]["output"]["attention_backend"] == ( + KRONOS2_XFORMERS_ATTENTION_BACKEND + ) + + monkeypatch.setenv("XFORMERS_DISABLED", "1") + pytorch = patch_output_fingerprint(spec, "fp32") + assert resolved_kronos2_attention_backend() == (KRONOS2_PYTORCH_ATTENTION_BACKEND) + assert pytorch["payload"]["output"]["attention_backend"] == ( + KRONOS2_PYTORCH_ATTENTION_BACKEND + ) + assert pytorch["digest"] != xformers["digest"] + + # Upstream treats presence, rather than truthiness, as the disable switch. + monkeypatch.setenv("XFORMERS_DISABLED", "") + assert resolved_kronos2_attention_backend() == (KRONOS2_PYTORCH_ATTENTION_BACKEND) + + +def test_resolved_backend_honours_an_already_imported_upstream_fallback(monkeypatch): + import raw2features.embedders.kronos2_embedder as module + + monkeypatch.delenv("XFORMERS_DISABLED", raising=False) + monkeypatch.setattr(module, "_allows_pinned_xformers", lambda: True) + monkeypatch.setitem( + sys.modules, + "dinov2.layers.attention", + SimpleNamespace(XFORMERS_AVAILABLE=False), + ) + + assert resolved_kronos2_attention_backend() == (KRONOS2_PYTORCH_ATTENTION_BACKEND) + + +def test_loaded_contract_uses_the_backend_recorded_during_model_load(monkeypatch): + import raw2features.embedders.kronos2_embedder as module + + monkeypatch.setattr(module, "_allows_pinned_xformers", lambda: False) + spec = get_spec("kronos2") + embedder = SimpleNamespace( + name="kronos2", + spec=spec, + embedding_dim=spec.embedding_dim, + _dtype=None, + _device="cuda", + _resolved_attention_backend=KRONOS2_XFORMERS_ATTENTION_BACKEND, + ) + + fingerprint = _loaded_model_contracts([embedder])["kronos2"]["output_fingerprint"] + assert fingerprint["payload"]["output"]["attention_backend"] == ( + KRONOS2_XFORMERS_ATTENTION_BACKEND + ) + + del embedder._resolved_attention_backend + with pytest.raises(RuntimeError, match="no resolved attention-backend"): + _loaded_model_contracts([embedder]) + + def test_forced_fallback_rejects_an_already_imported_xformers_backend(monkeypatch): import raw2features.embedders.kronos2_embedder as module @@ -866,6 +938,7 @@ def from_pretrained(source, **kwargs): ) ] assert emb._dtype is torch.float32 + assert emb._resolved_attention_backend == KRONOS2_PYTORCH_ATTENTION_BACKEND assert set(emb._metadata_index) == {"dapi", "draq5"} assert not hasattr(emb, "_registered_additional_markers_digest") assert not hasattr(emb, "_novel_registration_failed") From 9f09f2cb466f8f73c490c0f55880b50385972291 Mon Sep 17 00:00:00 2001 From: Craig Myles Date: Tue, 4 Aug 2026 22:21:11 +0100 Subject: [PATCH 4/4] Make KRONOS2 verification backend-portable --- CHANGELOG.md | 3 +- src/raw2features/embedders/fingerprint.py | 41 ++++++++++ src/raw2features/pipeline/receipt.py | 10 ++- src/raw2features/pipeline/runner.py | 18 ++++- tests/test_kronos2.py | 95 ++++++++++++++++++++++- 5 files changed, 156 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9a6bb4a..bcb6e80 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,7 +27,8 @@ layout, geometry-only `grid_hash`, and brightfield execution path are unchanged. (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 - produced by the two paths cannot be silently combined during resume. + 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. diff --git a/src/raw2features/embedders/fingerprint.py b/src/raw2features/embedders/fingerprint.py index fa4b5a0..2e84bf8 100644 --- a/src/raw2features/embedders/fingerprint.py +++ b/src/raw2features/embedders/fingerprint.py @@ -265,6 +265,47 @@ def output_fingerprints_equal(left: Any, right: Any) -> bool: ) +def output_fingerprints_compatible(left: Any, right: Any) -> bool: + """Compare output contracts, allowing only KRONOS2's approved backends. + + A completed KRONOS2 array remains valid when inspected on a machine that would + select the other supported attention implementation. The concrete backend stays + in the fingerprint; every other output-affecting field must remain identical. + """ + + if output_fingerprints_equal(left, right): + return True + if not valid_output_fingerprint(left) or not valid_output_fingerprint(right): + return False + + left_payload = deepcopy(dict(left["payload"])) + right_payload = deepcopy(dict(right["payload"])) + if any( + payload.get("kind") != "patch_features" + or payload.get("model") != "kronos2" + or not isinstance(payload.get("loader"), dict) + or payload["loader"].get("family") != "kronos2" + or not isinstance(payload.get("output"), dict) + for payload in (left_payload, right_payload) + ): + return False + + from .kronos2_embedder import ( + KRONOS2_PYTORCH_ATTENTION_BACKEND, + KRONOS2_XFORMERS_ATTENTION_BACKEND, + ) + + approved = ( + KRONOS2_PYTORCH_ATTENTION_BACKEND, + KRONOS2_XFORMERS_ATTENTION_BACKEND, + ) + left_backend = left_payload["output"].pop("attention_backend", None) + right_backend = right_payload["output"].pop("attention_backend", None) + if left_backend not in approved or right_backend not in approved: + return False + return left_payload == right_payload + + def fingerprint_digest(value: Any) -> str | None: """Return a validated record's digest, else ``None``.""" diff --git a/src/raw2features/pipeline/receipt.py b/src/raw2features/pipeline/receipt.py index e454e6f..4adaa46 100644 --- a/src/raw2features/pipeline/receipt.py +++ b/src/raw2features/pipeline/receipt.py @@ -220,6 +220,7 @@ def validate_model( return False if expected_fingerprint is not None: from raw2features.embedders.fingerprint import ( + output_fingerprints_compatible, output_fingerprints_equal, ) @@ -235,9 +236,14 @@ def validate_model( if isinstance(model_meta, dict) else None ) - if not output_fingerprints_equal(array_fingerprint, expected_fingerprint): + # Header and array must describe one exact producer. Once that invariant + # holds, a completed KRONOS2 output may be checked on a host whose approved + # attention backend differs from the producer recorded in the store. + if not output_fingerprints_equal(array_fingerprint, header_fingerprint): return False - if not output_fingerprints_equal(header_fingerprint, expected_fingerprint): + if not output_fingerprints_compatible( + array_fingerprint, expected_fingerprint + ): return False block = 8192 for s in range(0, n_patches, block): diff --git a/src/raw2features/pipeline/runner.py b/src/raw2features/pipeline/runner.py index c360d66..481dc97 100644 --- a/src/raw2features/pipeline/runner.py +++ b/src/raw2features/pipeline/runner.py @@ -31,6 +31,7 @@ from raw2features.embedders.fingerprint import ( expected_patch_outputs, make_output_fingerprint, + output_fingerprints_compatible, patch_output_fingerprint, resolved_patch_amp, ) @@ -1544,7 +1545,7 @@ def embedder_factory(device: str) -> list: ) ) run_contracts = {name: model_contracts[name] for name in models_to_do} - _assert_loaded_model_contracts( + run_contracts = _assert_loaded_model_contracts( run_embedders, run_contracts, cfg.resolved_channel_names, @@ -2566,8 +2567,8 @@ def _assert_loaded_model_contracts( selected_channels: list[dict[str, Any]] | None = None, explicit_selection: bool = False, native_model_params: dict[str, dict[str, Any]] | None = None, -) -> None: - """Fail before store mutation if loaded model copies differ from provenance.""" +) -> dict[str, dict]: + """Return authoritative loaded contracts after compatibility validation.""" actual = _loaded_model_contracts( embedders, @@ -2583,13 +2584,22 @@ def _assert_loaded_model_contracts( "Loaded model set does not match the output contract " f"(missing={missing}, unexpected={extra})." ) - mismatched = [name for name in expected if actual[name] != expected[name]] + mismatched = [ + name + for name in expected + if actual[name].get("embedding_dim") != expected[name].get("embedding_dim") + or not output_fingerprints_compatible( + actual[name].get("output_fingerprint"), + expected[name].get("output_fingerprint"), + ) + ] if mismatched: raise ValueError( "Loaded model contract differs from the requested/persisted contract for " f"{mismatched}. Check effective AMP, preprocessing, dimensions, and " "weights." ) + return actual def load_embedders(cfg: RunConfig, device: str | None = None) -> list: diff --git a/tests/test_kronos2.py b/tests/test_kronos2.py index bfc9cea..5e213a3 100644 --- a/tests/test_kronos2.py +++ b/tests/test_kronos2.py @@ -13,6 +13,8 @@ import pytest from raw2features.embedders.fingerprint import ( + make_output_fingerprint, + output_fingerprints_compatible, patch_output_fingerprint, resolved_patch_amp, ) @@ -40,8 +42,10 @@ parse_kronos2_additional_markers, ) from raw2features.embedders.model_registry import get_spec +from raw2features.pipeline.receipt import validate_model from raw2features.pipeline.runner import ( RunConfig, + _assert_loaded_model_contracts, _loaded_model_contracts, expected_model_contracts, ) @@ -789,10 +793,83 @@ def test_resolved_backend_honours_an_already_imported_upstream_fallback(monkeypa assert resolved_kronos2_attention_backend() == (KRONOS2_PYTORCH_ATTENTION_BACKEND) +def test_backend_compatibility_is_narrow_and_keeps_store_provenance_exact(tmp_path): + import zarr + + spec = get_spec("kronos2") + xformers = patch_output_fingerprint( + spec, + "fp32", + resolved_attention_backend=KRONOS2_XFORMERS_ATTENTION_BACKEND, + ) + pytorch = patch_output_fingerprint( + spec, + "fp32", + resolved_attention_backend=KRONOS2_PYTORCH_ATTENTION_BACKEND, + ) + assert output_fingerprints_compatible(xformers, pytorch) + + changed_payload = json.loads(json.dumps(pytorch["payload"])) + changed_payload["checkpoint"]["weights_revision"] = "different" + changed = make_output_fingerprint(changed_payload) + assert not output_fingerprints_compatible(xformers, changed) + + missing_backend_payload = json.loads(json.dumps(pytorch["payload"])) + del missing_backend_payload["output"]["attention_backend"] + missing_backend = make_output_fingerprint(missing_backend_payload) + assert not output_fingerprints_compatible(xformers, missing_backend) + + unapproved_backend_payload = json.loads(json.dumps(pytorch["payload"])) + unapproved_backend_payload["output"]["attention_backend"] = { + "implementation": "unvalidated_attention", + "package": "unknown", + } + unapproved_backend = make_output_fingerprint(unapproved_backend_payload) + assert not output_fingerprints_compatible(xformers, unapproved_backend) + + group = zarr.open_group(str(tmp_path / "kronos2.zarr"), mode="w", zarr_format=2) + features = group.create_group("features") + array = features.create_array("kronos2", shape=(2, 768), dtype="float32") + array[:] = 1 + array.attrs["output_fingerprint"] = xformers + group.attrs["raw2features"] = { + "models": {"kronos2": {"output_fingerprint": xformers}} + } + + # A complete xFormers output remains valid when checked on a PyTorch-fallback + # host, while the two persisted copies must still identify the same producer. + assert validate_model( + group, + "kronos2", + 2, + expected_dim=768, + expected_fingerprint=pytorch, + ) + group.attrs["raw2features"] = { + "models": {"kronos2": {"output_fingerprint": pytorch}} + } + assert not validate_model( + group, + "kronos2", + 2, + expected_dim=768, + expected_fingerprint=pytorch, + ) + del array.attrs["output_fingerprint"] + assert not validate_model( + group, + "kronos2", + 2, + expected_dim=768, + expected_fingerprint=pytorch, + ) + + def test_loaded_contract_uses_the_backend_recorded_during_model_load(monkeypatch): import raw2features.embedders.kronos2_embedder as module - monkeypatch.setattr(module, "_allows_pinned_xformers", lambda: False) + monkeypatch.delenv("XFORMERS_DISABLED", raising=False) + monkeypatch.setattr(module, "_allows_pinned_xformers", lambda: True) spec = get_spec("kronos2") embedder = SimpleNamespace( name="kronos2", @@ -800,14 +877,24 @@ def test_loaded_contract_uses_the_backend_recorded_during_model_load(monkeypatch embedding_dim=spec.embedding_dim, _dtype=None, _device="cuda", - _resolved_attention_backend=KRONOS2_XFORMERS_ATTENTION_BACKEND, + _resolved_attention_backend=KRONOS2_PYTORCH_ATTENTION_BACKEND, ) - fingerprint = _loaded_model_contracts([embedder])["kronos2"]["output_fingerprint"] - assert fingerprint["payload"]["output"]["attention_backend"] == ( + expected = expected_model_contracts( + RunConfig(models=["kronos2"], no_seg=True, patch_px=256) + ) + expected_fingerprint = expected["kronos2"]["output_fingerprint"] + assert expected_fingerprint["payload"]["output"]["attention_backend"] == ( KRONOS2_XFORMERS_ATTENTION_BACKEND ) + actual = _assert_loaded_model_contracts([embedder], expected) + fingerprint = actual["kronos2"]["output_fingerprint"] + assert fingerprint["payload"]["output"]["attention_backend"] == ( + KRONOS2_PYTORCH_ATTENTION_BACKEND + ) + assert output_fingerprints_compatible(fingerprint, expected_fingerprint) + del embedder._resolved_attention_backend with pytest.raises(RuntimeError, match="no resolved attention-backend"): _loaded_model_contracts([embedder])