Skip to content

Add pinned KRONOS2 native multiplex embeddings - #19

Draft
CraigMyles wants to merge 1 commit into
mainfrom
feature/kronos2
Draft

Add pinned KRONOS2 native multiplex embeddings#19
CraigMyles wants to merge 1 commit into
mainfrom
feature/kronos2

Conversation

@CraigMyles

Copy link
Copy Markdown
Owner

Summary

  • add kronos2 as a native multiplex patch encoder using the immutable MahmoodLab/KRONOS2 Hugging Face revision and SHA-verified weights and marker metadata
  • support ordered native marker selection and optional novel-marker registration through the upstream BioLinkBERT path, with the conditional text model pinned and every required artifact verified
  • bind the effective panel, physical marker indices, normalization, forward contract, and optional registration recipe into output fingerprints and store provenance
  • prepare the focused v0.2.1 model update without changing the v0.1 store layout, geometry-only grid_hash, or ordinary brightfield execution

Model loading and execution

  • download only the pinned KRONOS2 files required for inference into a
    raw2features-managed cache; reject unexpected files before executing the model's
    custom code
  • construct the model entirely from that local pinned copy, then remove the temporary
    Python import path added by the upstream loader
  • use the published fp32 inference path and fixed batches of eight; short final batches
    are temporarily padded by repeating their last patch, with the extra outputs
    discarded
  • accept square patch sizes divisible by 16, defaulting to 256 px; reject
    torch.compile because compiled execution has not been validated against the
    reference implementation
  • use the 288 released marker entries without downloading BioLinkBERT; the additional
    text model is downloaded only when a user explicitly registers novel markers
  • do not silently treat Hoechst or DNA1/DNA2 as DAPI because KRONOS2 applies
    marker-specific normalization; excluded nuclear channels remain available for
    segmentation, while explicit novel-marker registration is supported
  • keep native marker-aware models such as KRONOS2 separate from RGB encoders within a
    single command; compatible runs can still add their outputs to the same store

Validation

  • 678 local non-slow, non-network tests passed; 105 dependency-gated tests skipped and 68 slow/network tests deselected
  • the loopback authenticated-HTTP reader test passed separately outside the filesystem sandbox
  • Ruff, git diff --check, uv lock --check, wheel/sdist builds, and Twine metadata checks passed
  • real-weight L40S validation used Python 3.12, Torch 2.6.0+cu124, Transformers 4.56.0, timm 1.0.19, and xFormers 0.0.29.post3
  • verified the pinned KRONOS2 weight and marker-table hashes before loading; the
    raw2features-managed cache contained only its seven required model entries plus
    Hugging Face's cache metadata, and xFormers was active
  • on a real uint16 eight-channel IMC OME-Zarr, the published-marker path produced finite deterministic (9, 768) fp32 embeddings and matched the upstream forward bit-for-bit for the first full batch
  • the optional pinned BioLinkBERT path registered a novel marker and produced finite (2, 768) fp32 embeddings
  • marker reorder, exclusion, and novel registration produced distinct output fingerprints
  • an end-to-end CLI run wrote 16 finite (768,) KRONOS2 embeddings with the expected marker-sensitive fingerprint

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant