diff --git a/CHANGELOG.md b/CHANGELOG.md index 32750ee..bcb6e80 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,42 @@ Notable changes to raw2features, newest first. This project follows [Semantic Versioning](https://semver.org). +## [0.2.1] - Unreleased + +This focused model update adds KRONOS2 to the existing native multiplex path. The store +layout, geometry-only `grid_hash`, and brightfield execution path are unchanged. + +### Models and multiplex inputs + +- Add the gated `kronos2` native multiplex patch encoder: a pinned, SHA-verified + marker-aware ViT-B/16 that produces a 768-dimensional CLS vector from ordered marker + stacks. The released metadata contains 288 usable marker entries, 268 of them + marked as pretraining markers. This path runs in fp32 and does not download its + BioLinkBERT text encoder. Its recorded forward contract uses fixed eight-sample + microbatches, including deterministic padding of a short final microbatch. The + optional extra installs the authors' pinned x86-64 Linux reference stack. Only the + required pinned custom-code modules enter an app-owned runtime snapshot, and the + upstream-added import path is removed after model construction. +- Use DRAQ5 as the nuclear-segmentation fallback when a multiplex panel has no + recognized DAPI, Hoechst, or DNA stain. Existing panels containing DAPI retain their + DAPI binding, and repeated DRAQ5 acquisitions are averaged with their physical indices + recorded in grid identity. +- Install the pinned xFormers reference kernel only where its release provides wheels + (x86-64 Linux, CPython 3.11/3.12); Python 3.13 and other platforms use KRONOS2's + enforced upstream PyTorch attention fallback instead of an unvalidated xFormers build. + Record the resolved attention backend in each model-output fingerprint so outputs + remain self-describing across hosts; completed arrays from either approved path remain + valid, while exact header/array agreement prevents partial writes being accepted. +- Extend repeated `--marker` selection and ordering to native multiplex encoders. The + resolved physical indices and order are included in output fingerprints and panel + provenance without changing patch-grid identity. +- Add optional `--kronos2-additional-markers` registration for novel markers. The + canonical consumed CSV fields, intensity statistics, and pinned BioLinkBERT dependency + are bound into the model contract; BioLinkBERT is fetched only when this option is used. +- Keep native multiplex and ordinary RGB encoders in separate requests so each receives + the correct source-channel view. Additive runs can still write their outputs into the + same store when their grid geometry matches. + ## [0.2.0] - 2026-07-23 This release extends named-channel multiplex support with a `channelwise` strategy, adds @@ -130,3 +166,4 @@ What 0.1.0 provides: [0.1.0]: https://github.com/CraigMyles/raw2features/releases/tag/v0.1.0 [0.2.0]: https://github.com/CraigMyles/raw2features/compare/v0.1.0...v0.2.0 +[0.2.1]: https://github.com/CraigMyles/raw2features/compare/v0.2.0...HEAD diff --git a/CITATION.cff b/CITATION.cff index c294b44..d123f87 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -20,5 +20,4 @@ keywords: - foundation models - feature extraction license: MIT -version: "0.2.0" -date-released: 2026-07-23 +version: "0.2.1" diff --git a/README.md b/README.md index b53053f..c7229c6 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,25 @@ The same pattern covers the other gated encoders - mostly slide encoders (e.g. ` or Drive-hosted weights). Each model's exact install is in its [`MODELS.md`](docs/MODELS.md) row and the matching extra's comment in `pyproject.toml`. +**KRONOS2.** The gated KRONOS2 repository is loaded from a pinned Hugging Face snapshot, +so it needs no separate git-package install: + +```bash +pip install "raw2features[all,kronos2]" +hf auth login +``` + +Request access to [`MahmoodLab/KRONOS2`](https://huggingface.co/MahmoodLab/KRONOS2) +with the institutional account required by its gate. On x86-64 Linux with CPython 3.11 +or 3.12, the extra installs the authors' pinned Torch/xFormers/Transformers/timm +reference stack; raw2features validates this path on Python 3.12. Python 3.13 and other +platforms use the upstream PyTorch attention fallback, which is slower and may use more +memory. The KRONOS2 extra is not available on Python 3.14 because the pinned Torch +release does not support it. Normal use of the released marker vocabulary does not +download BioLinkBERT; that conditional download occurs only when a complete novel-marker +CSV is supplied. See +[MODALITIES.md](docs/MODALITIES.md) for marker selection and registration. + **Development** (from a clone, with [uv](https://docs.astral.sh/uv/)): ```bash diff --git a/docs/MODALITIES.md b/docs/MODALITIES.md index 3e13f1c..ef9c1d8 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,55 @@ full CLI contract are documented in [usage.md](usage.md#rgb-encoders-on-named-ch ## Native multiplex encoders Native multiplex models use the same positional-panel plumbing but consume the marker -stack directly. The current built-in example, `kronos`, maps effective source names to -its pinned marker vocabulary, records the kept/dropped physical mapping, and binds the -complete effective panel into its output fingerprint before resume. Installation and -model-specific license/access details are in [MODELS.md](MODELS.md) and -[MODEL_LICENSES.md](MODEL_LICENSES.md). +stack directly. `kronos` and `kronos2` map effective source names to their pinned marker +metadata, record the physical mapping, and bind the selected indices and order into the +output fingerprint before resume. Repeated `--marker` options select and order a subset; +without them, the source channels are offered to the model in physical order. +Installation and model-specific license/access details are in +[MODELS.md](MODELS.md) and [MODEL_LICENSES.md](MODEL_LICENSES.md). ```bash raw2features embed SLIDE.ome.zarr OUT -m kronos --mpp 0.5 +raw2features embed SLIDE.ome.zarr OUT -m kronos2 \ + --marker CD3 --marker CD8 --marker DAPI ``` Use `--channel-names-file` when the source lacks a complete panel, or `--no-seg` when the entire image should be tiled without a nuclear mask. +Nuclear masking prefers recognized DAPI, Hoechst, or DNA channels. If none is present, +it uses DRAQ5 as a fallback. Repeated acquisitions of the selected stain are averaged; +DAPI takes precedence over DRAQ5 rather than mixing the two dyes. The resolved physical +indices and combination are recorded in grid identity. + +KRONOS2 performs exact separator-insensitive marker matching and does not reuse +KRONOSv1's biological aliases. Its released metadata contains 288 usable entries, of +which 268 are marked as pretraining markers. Their text vectors are already stored in +the checkpoint, so BioLinkBERT is neither downloaded nor run. + +The released KRONOS2 vocabulary includes DAPI and DRAQ5, but not Hoechst or +DNA1/DNA2 labels found in many CODEX panels. raw2features does not silently map those +stains to DAPI. A Hoechst channel may be excluded from the encoder with an explicit +`--marker` list while still being used by nuclear segmentation, which always sees the +complete source panel. To include it in the KRONOS2 marker stack, register it as a novel +marker with statistics computed for the relevant dataset. + +For a marker outside the released metadata, pass a complete CSV through +`--kronos2-additional-markers`. Each row requires `marker_name`, `marker_full_name`, +`compartment`, `family`, `family_desc`, `mean`, and `std`. The statistics describe the +marker after KRONOS2's published dtype scaling (uint8 / 255, wider unsigned integers / +65535, or floating values / 400) and are supplied by the user; signed integers are +rejected and raw2features does not infer cohort statistics. Novel names must not collide +with the pinned vocabulary, while `compartment` and `family` must reuse its categories +exactly. This optional path downloads the separately pinned BioLinkBERT model to create +the marker text vector. The CSV content, statistics, and text-model identity become part +of the output fingerprint. + +Following the published inference example, raw2features passes DRAQ5 as +`preferred_dapi` when DAPI is absent. The upstream preprocessor therefore applies its +DAPI normalization statistics to that selected DRAQ5 channel. This policy is recorded +in provenance and is not configurable in v0.2.1. + ## Extending multiplex support The `raw2features.multiplex_strategies` entry point separates panel/config preparation diff --git a/docs/MODELS.md b/docs/MODELS.md index 7113904..ae8f519 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,31 @@ back to **1.0 µm/px** unless run alongside a model that supplies one or given a SHA-256-verified safetensors file and does not execute the repository's remote code or construct its unused BERT text tower. - `kronos` - needs the `[kronos]` extra; multiplex only (see `MODALITIES.md`). +- `kronos2` - needs the `[kronos2]` extra and accepted institutional access to the gated + [`MahmoodLab/KRONOS2`](https://huggingface.co/MahmoodLab/KRONOS2) repository. It loads + an allowlisted pinned custom-code snapshot into an app-owned runtime directory, verifies + the weights and marker table, and removes the upstream-added snapshot import path after + construction. It runs in fp32 and returns the 768-dimensional CLS token. The default + patch is 256 px; + an explicit alternative must be a multiple of the ViT patch size, 16. Marker matching + is exact after separator-insensitive normalization and does not apply KRONOSv1 aliases. + Upstream recommends batches of at least eight for reproducibility. raw2features fixes + its effective forward microbatch at eight and pads a shorter final microbatch before + slicing the output back to its original length. This runtime contract is recorded; + bitwise equality across different hardware is not promised. On x86-64 Linux with + CPython 3.11 or 3.12, the `[kronos2]` extra installs the authors' pinned Torch, + xFormers, Transformers, and timm reference stack; this path is validated with Python + 3.12. Python 3.13 and other platforms use the upstream PyTorch attention fallback, + which is slower and may use more memory. The pinned Torch release does not support + Python 3.14. `--compile` is not supported for KRONOS2 in v0.2.1. + The ordinary 288-marker path uses text embeddings already stored in the checkpoint + and never downloads BioLinkBERT. + The vocabulary includes DAPI and DRAQ5 but not Hoechst or DNA1/DNA2 labels; exclude + those channels explicitly or register them with dataset-specific statistics. Following + the upstream example, a selected DRAQ5 channel uses DAPI normalization statistics when + DAPI is absent; this recorded policy is not configurable in v0.2.1. + Supplying `--kronos2-additional-markers` enables the separately pinned BioLinkBERT path + for complete user-defined marker metadata and statistics. - `ctranspath` - custom `ConvStem` patch-embed (`embedders/convstem.py`), loaded via the 1aurent modern-timm mirror (no pinned timm fork needed). - `gpfm`, `retccl`, `sp22m`, `hipt`, `lunit_bt/mocov2/swav` - built from a base timm arch + 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..cb26c41 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,77 @@ 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. 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 +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 +420,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 +509,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..1c7eff1 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 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' 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", + "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..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``.""" @@ -297,6 +338,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 +419,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 +457,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, @@ -435,7 +513,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: @@ -529,12 +612,38 @@ 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, + 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) 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..733f163 --- /dev/null +++ b/src/raw2features/embedders/kronos2_embedder.py @@ -0,0 +1,987 @@ +"""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 +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. +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 platform +import sys +import tempfile +import unicodedata +from importlib import metadata as importlib_metadata +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_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", + "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", + "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", + "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 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 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.""" + + 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.""" + + @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 the directory is not left at the front of the caller's + # import search path. + 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() + os.environ["XFORMERS_DISABLED"] = "1" + snapshot_path_count = sys.path.count(snapshot) + try: + model = AutoModel.from_pretrained( + snapshot, + trust_remote_code=True, + local_files_only=True, + device=device, + ) + finally: + 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 + 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 + # 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..4b71142 100644 --- a/src/raw2features/embedders/registry.yaml +++ b/src/raw2features/embedders/registry.yaml @@ -1011,6 +1011,43 @@ 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. 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` + 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/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 600d635..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, ) @@ -106,6 +107,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 +272,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 +290,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 +327,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 +341,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 +416,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 +446,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 +521,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 +555,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 +563,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 +625,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 +686,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 +728,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 +785,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 +798,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 +896,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 +964,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 +986,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 +1340,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 @@ -1297,8 +1545,13 @@ 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_contracts = _assert_loaded_model_contracts( + 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 +1596,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 +2035,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 +2146,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 +2339,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 +2350,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 +2365,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 +2382,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 +2485,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 +2514,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.""" @@ -2216,15 +2531,31 @@ 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, [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 +2563,20 @@ def _assert_loaded_model_contracts( embedders: list, expected: dict[str, dict], channel_names: list[str] | None = None, -) -> None: - """Fail before store mutation if loaded model copies differ from provenance.""" + *, + 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]: + """Return authoritative loaded contracts after compatibility validation.""" - 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)) @@ -2243,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: @@ -2367,7 +2717,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"} @@ -2579,6 +2930,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 +3036,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 +3067,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 +3112,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 +3220,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/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"(? 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["forward"]["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", + }, + } + 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", + "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"): + 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) + + +@pytest.mark.parametrize( + ("system", "machine", "implementation", "python_version", "version", "expected"), + [ + ("Linux", "x86_64", "CPython", (3, 11), KRONOS2_XFORMERS_VERSION, True), + ("Linux", "AMD64", "CPython", (3, 12), KRONOS2_XFORMERS_VERSION, True), + ("Linux", "x86_64", "CPython", (3, 13), KRONOS2_XFORMERS_VERSION, False), + ("Linux", "aarch64", "CPython", (3, 12), KRONOS2_XFORMERS_VERSION, False), + ("Darwin", "x86_64", "CPython", (3, 12), KRONOS2_XFORMERS_VERSION, False), + ("Linux", "x86_64", "PyPy", (3, 12), KRONOS2_XFORMERS_VERSION, False), + ("Linux", "x86_64", "CPython", (3, 12), "0.0.30", False), + ], +) +def test_only_the_pinned_reference_xformers_environment_is_allowed( + monkeypatch, system, machine, implementation, python_version, version, expected +): + import raw2features.embedders.kronos2_embedder as module + + monkeypatch.setattr(module.importlib_metadata, "version", lambda _name: version) + + assert ( + _allows_pinned_xformers( + system=system, + machine=machine, + implementation=implementation, + python_version=python_version, + ) + is expected + ) + + +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_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.delenv("XFORMERS_DISABLED", raising=False) + monkeypatch.setattr(module, "_allows_pinned_xformers", lambda: True) + spec = get_spec("kronos2") + embedder = SimpleNamespace( + name="kronos2", + spec=spec, + embedding_dim=spec.embedding_dim, + _dtype=None, + _device="cuda", + _resolved_attention_backend=KRONOS2_PYTORCH_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]) + + +def test_forced_fallback_rejects_an_already_imported_xformers_backend(monkeypatch): + import raw2features.embedders.kronos2_embedder as module + + monkeypatch.setitem( + sys.modules, + "dinov2.layers.attention", + SimpleNamespace(XFORMERS_AVAILABLE=True), + ) + with pytest.raises(RuntimeError, match="already-imported DINOv2"): + module._assert_forced_fallback_is_not_preimported_with_xformers() + + +@pytest.mark.parametrize("snapshot_preexisting", [False, True]) +def test_loader_verifies_both_artifacts_then_loads_only_the_local_snapshot( + monkeypatch, tmp_path, snapshot_preexisting +): + 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 = [] + backend_env = [] + 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) + monkeypatch.setattr(module, "_allows_pinned_xformers", lambda: False) + monkeypatch.delenv("XFORMERS_DISABLED", raising=False) + 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)) + backend_env.append(module.os.environ.get("XFORMERS_DISABLED")) + 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,) + if snapshot_preexisting: + monkeypatch.setattr(sys, "path", [*sys.path, str(tmp_path)]) + original_sys_path = list(sys.path) + 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 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") + assert not hasattr(emb, "_raw2features_multiplex_source_indices") + assert sys.path == original_sys_path + assert backend_env == ["1"] + assert "XFORMERS_DISABLED" not in module.os.environ + + +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_multiplex.py b/tests/test_multiplex.py index 9f1266e..6b77711 100644 --- a/tests/test_multiplex.py +++ b/tests/test_multiplex.py @@ -111,6 +111,27 @@ def test_nuclear_segmenter_thresholds_dapi(synthetic_multiplex_ngff): assert set(np.unique(tm.mask)).issubset({0.0, 1.0}) +def test_nuclear_segmenter_thresholds_draq5_as_a_fallback( + synthetic_multiplex_ngff, +): + import zarr + + from raw2features.segmenters.nuclear import NuclearSegmenter + + root = zarr.open_group(synthetic_multiplex_ngff, mode="r+") + labels = ["DRAQ5", "CD3", "CD8", "CD20", "FOXP3"] + root.attrs["omero"] = { + "channels": [{"label": label, "active": True} for label in labels] + } + + with OmeZarrReader(synthetic_multiplex_ngff) as reader: + assert NuclearSegmenter()._nuclear_indices(reader.channel_names) == [0] + mask = NuclearSegmenter(seg_mpp=2.0).segment(reader) + + assert mask.mask.ndim == 2 and mask.mask.dtype == np.float32 + assert set(np.unique(mask.mask)).issubset({0.0, 1.0}) + + def test_nuclear_segmenter_errors_without_a_nuclear_channel(synthetic_multiplex_ngff): from raw2features.segmenters.nuclear import NuclearSegmenter @@ -191,6 +212,46 @@ def test_nuclear_segmenter_combines_metal_prefixed_dna_pair_and_repeated_dapi(): ] +@pytest.mark.parametrize( + "name", + ["DRAQ5", "DRAQ-5", "DRAQ 5", "191Ir_DRAQ5", "DRAQ5(Ir191)"], +) +def test_nuclear_segmenter_recognizes_bounded_draq5_variants(name): + from raw2features.segmenters.nuclear import NuclearSegmenter + + assert NuclearSegmenter()._nuclear_indices(["CD3", name]) == [1] + + +@pytest.mark.parametrize("name", ["DRAQ7", "antiDRAQ5", "DRAQ51"]) +def test_nuclear_segmenter_rejects_draq5_false_positives(name): + from raw2features.segmenters.nuclear import NuclearSegmenter + + with pytest.raises(ValueError, match="no nuclear channel"): + NuclearSegmenter()._nuclear_indices([name, "CD3"]) + + +def test_nuclear_segmenter_uses_draq5_only_when_primary_stain_is_absent(): + from raw2features.segmenters.nuclear import NuclearSegmenter + + segmenter = NuclearSegmenter() + assert segmenter._nuclear_indices(["DRAQ5", "DAPI", "CD3"]) == [1] + assert segmenter._nuclear_indices(["DAPI", "DRAQ5", "CD3"]) == [0] + assert segmenter._nuclear_indices(["DRAQ5", "Hoechst1", "CD3"]) == [1] + + +def test_nuclear_segmenter_combines_repeated_draq5_without_uint16_overflow(): + from raw2features.segmenters.nuclear import NuclearSegmenter + + segmenter = NuclearSegmenter() + assert segmenter._nuclear_indices(["CD3", "DRAQ5", "DRAQ-5"]) == [1, 2] + block = np.zeros((2, 2, 3), dtype=np.uint16) + block[..., 1] = np.uint16(60_000) + block[..., 2] = np.uint16(50_000) + combined = segmenter._combine_nuclear_channels(block, [1, 2]) + assert combined.dtype == np.float32 + assert np.all(combined == 55_000.0) + + @pytest.mark.parametrize( "names", [ 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_release_metadata.py b/tests/test_release_metadata.py index 2c38105..754048c 100644 --- a/tests/test_release_metadata.py +++ b/tests/test_release_metadata.py @@ -6,6 +6,7 @@ from pathlib import Path import yaml +from packaging.requirements import Requirement import raw2features @@ -19,4 +20,44 @@ def test_release_versions_match() -> 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/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..af8a22a 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 = "python_full_version < '3.13' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' 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 = "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", "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 = "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 = [ + { 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"