| """Four-modality embedding model: H&E, mIF (protein), RNA (spot expression) and text in ONE 1536-d space. |
| |
| Configuration EJ-ad+MV (RESULTS.md, "MedVTok as H&E image information + adapter-only joint training"): |
| stage 1 every modality encoder gets its own head into the frozen UNI2-h space (InfoNCE against UNI2 on REAL pairs: |
| ORION H&E-mIF, HEST H&E-RNA, TCGA/PathGen H&E-text) |
| mIF : VirTues (16-marker ORION panel) -> MLP |
| RNA : hvg (log1p CP10k of 2000 HEST HVGs), Tahoe-x1 70M, SE-600M, AIDO.Cell-100M -> one linear head each |
| text : KEEP, BiomedCLIP (text tower), Qwen3-Embedding-0.6B -> one linear head each |
| RNA / text = late fusion: z = normalise(sum_e normalise(head_e(x_e))) |
| stage 2 all mIF / RNA / text heads frozen; only a zero-initialised residual H&E adapter (+ temperatures) trains: |
| H&E = UNI2 + A([UNI2, c * MedVTok]) |
| so mIF / RNA / text embeddings (and every emergent mIF<->RNA<->text pairing) are exactly the stage-1 ones. |
| Checkpoints are written by 52_final_model.py (models/<name>/model.pt + meta.json). |
| |
| Usage (python: .venv/bin/python -I, run from anywhere; GPU): |
| from embed import Embedder |
| E = Embedder("models/ej_ad_mv") # release model |
| z = E.he(images) # H&E tiles: list of PIL images / (N, H, W, 3) uint8; ~0.5 um/px, 112-166 um field |
| z = E.text(["Tumor cells with ...", ...]) |
| z = E.rna(counts, genes, zscore="train") # raw counts (spots x genes; numpy or scipy.sparse), HGNC symbols |
| z = E.mif(tiles, slide) # (N, 16, H, W) 8-bit-scale mIF tiles, channels = MIF_MARKERS, + slide id per tile |
| sim = z_a @ z_b.T # every output: L2-normalised (N, 1536) float32 torch tensor on CPU |
| Feature-level entry points (precomputed encoder features, as in cache/): |
| E.he_from_feats(uni2, medvtok), E.mif_from_feats(virtues), E.rna_from_feats({enc: X}), E.text_from_feats({enc: X}) |
| Raw-input pre-processing is exactly the one used to build the training features (functions imported from |
| 04b_ / 14_ / 43_ / 45_ / 45b_ / 50_ logic), see each method. Notes: |
| * RNA zscore="train" uses the HEST-train feature statistics; zscore="self" z-scores each FM feature over the given |
| spots (what the external DLPFC test does; use for a new dataset of >= a few hundred spots). |
| * mIF standardisation (per-slide 0.99-quantile clip, log1p, mean / std, blur) is computed over the tiles given in |
| one call, so embed a whole dataset (or at least whole slides) at once. |
| * Licences of the frozen encoders: UNI2-h (MahmoodLab, research), MedVTok (MIT; aligned to CONCH v1.5, CC-BY-NC-ND), |
| SE-600M and AIDO.Cell (non-commercial), Tahoe-x1 (Apache-2.0), VirTues, KEEP, BiomedCLIP (MIT), Qwen3 (Apache-2.0). |
| """ |
| import importlib.util |
| import os |
| import sys |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| ROOT = os.path.dirname(os.path.abspath(__file__)) |
| if ROOT not in sys.path: |
| sys.path.insert(0, ROOT) |
|
|
| DIM, MV_DIM = 1536, 1152 |
| RNA_E, TXT_E = ["hvg", "tahoex1", "se600m", "aidocell"], ["keep", "biomedclip", "qwen3-0.6b"] |
| MIF_MARKERS = ["Hoechst", "CD31", "CD45", "CD68", "CD4", "FOXP3", "CD8a", "CD45RO", "CD20", "PD-L1", "CD3e", "CD163", |
| "E-cadherin", "Ki67", "Pan-CK", "SMA"] |
| UNI2_LOCAL = "/diskarray/home/tl688/reg2026/model_alldata/uni2-h.bin" |
|
|
|
|
| def uni2_weights(): |
| """UNI2-h state dict path: $UNI2_PATH, else the local copy, else the gated HF repo MahmoodLab/UNI2-h""" |
| p = os.environ.get("UNI2_PATH", UNI2_LOCAL) |
| if os.path.exists(p): |
| return p |
| from huggingface_hub import hf_hub_download |
| return hf_hub_download("MahmoodLab/UNI2-h", "pytorch_model.bin") |
| nz = lambda x: F.normalize(x, dim=-1) |
|
|
|
|
| def mlp(din, dout=DIM, hid=512): |
| return nn.Sequential(nn.Linear(din, hid), nn.GELU(), nn.Dropout(0.1), nn.Linear(hid, hid), nn.GELU(), |
| nn.LayerNorm(hid), nn.Linear(hid, dout)) |
|
|
|
|
| class Bind4(nn.Module): |
| """Heads + H&E adapter. Inputs are raw encoder features; normalisation statistics live in buffers.""" |
|
|
| def __init__(self, dims, adapter="cat", n_adapters=1): |
| """dims: {'mif': 512, 'rna:hvg': 2000, ..., 'text:keep': 768, ...}; adapter None / 'u' / 'cat'; |
| n_adapters > 1: seed ensemble, H&E = UNI2 + mean_s A_s(.)""" |
| super().__init__() |
| self.keys = list(dims) |
| self.heads = nn.ModuleDict({self._k(k): (mlp(d) if k == "mif" else nn.Linear(d, DIM)) for k, d in dims.items()}) |
| self.logit = nn.ParameterDict({self._k(k): nn.Parameter(torch.tensor(np.log(1 / 0.07), dtype=torch.float32)) |
| for k in dims}) |
| for k, d in dims.items(): |
| self.register_buffer(f"mu_{self._k(k)}", torch.zeros(d)) |
| self.register_buffer(f"sd_{self._k(k)}", torch.ones(d)) |
| self.register_buffer(f"zs_{self._k(k)}", torch.tensor(False)) |
| self.adapter = adapter |
| self.register_buffer("c_mv", torch.tensor(1.0)) |
| self.ah = None |
| if adapter: |
| self.ah = nn.ModuleList([self.new_adapter() for _ in range(n_adapters)]) |
|
|
| def new_adapter(self): |
| """zero-initialised output -> starts exactly at H&E = UNI2""" |
| a = nn.Sequential(nn.Linear(DIM + (MV_DIM if self.adapter == "cat" else 0), 512), *mlp(DIM)[1:]) |
| nn.init.zeros_(a[-1].weight); nn.init.zeros_(a[-1].bias) |
| return a |
|
|
| @staticmethod |
| def _k(k): |
| return k.replace(":", "__").replace(".", "_").replace("-", "_") |
|
|
| def h(self, k): |
| return self.heads[self._k(k)] |
|
|
| def lg(self, k): |
| return self.logit[self._k(k)].exp().clamp(max=100) |
|
|
| def norm(self, k, x): |
| return (x - getattr(self, f"mu_{self._k(k)}")) / getattr(self, f"sd_{self._k(k)}") |
|
|
| def set_norm(self, k, mu, sd): |
| """x -> (x - mu) / sd; sd already includes the epsilon of the training code""" |
| getattr(self, f"mu_{self._k(k)}").copy_(mu); getattr(self, f"sd_{self._k(k)}").copy_(sd) |
| getattr(self, f"zs_{self._k(k)}").fill_(True) |
|
|
| |
| def he(self, uni2, medvtok=None): |
| if self.ah is None: |
| return uni2 |
| inp = uni2 if self.adapter == "u" else torch.cat([uni2, self.c_mv * medvtok], -1) |
| return uni2 + sum(a(inp) for a in self.ah) / len(self.ah) |
|
|
| def mif(self, x): |
| return self.h("mif")(self.norm("mif", x)) |
|
|
| def fuse(self, mod, X, zscore="train"): |
| """X: {enc: raw features}; zscore 'train' (stored stats) / 'self' (stats of X itself, raw encoders untouched)""" |
| z = {} |
| for e, x in X.items(): |
| k = f"{mod}:{e}" |
| raw = not bool(getattr(self, f"zs_{self._k(k)}")) |
| x = x if raw else (x - x.mean(0)) / (x.std(0) + 1e-6) if zscore == "self" else self.norm(k, x) |
| z[e] = nz(self.h(k)(x)) |
| return nz(sum(z.values())), z |
|
|
|
|
| def _load_script(fname, argv=None): |
| """import one of the numbered project scripts as a module (their CLI code is under __main__ / argv-guarded)""" |
| old = sys.argv |
| if argv is not None: |
| sys.argv = argv |
| try: |
| spec = importlib.util.spec_from_file_location("m_" + fname.split(".")[0], os.path.join(ROOT, fname)) |
| m = importlib.util.module_from_spec(spec); spec.loader.exec_module(m) |
| finally: |
| sys.argv = old |
| return m |
|
|
|
|
| def _as_uint8_list(images): |
| from PIL import Image, PngImagePlugin |
| PngImagePlugin.MAX_TEXT_CHUNK = max(PngImagePlugin.MAX_TEXT_CHUNK, 64 * 1024 * 1024) |
| out = [] |
| for im in images: |
| if isinstance(im, (str, os.PathLike)): |
| im = Image.open(im) |
| out.append(im.convert("RGB") if isinstance(im, Image.Image) else Image.fromarray(np.asarray(im, np.uint8))) |
| return out |
|
|
|
|
| class Embedder: |
| def __init__(self, path=os.path.join(ROOT, "models", "ej_ad_mv"), device="cuda"): |
| import json |
| self.dev = device |
| ck = torch.load(os.path.join(path, "model.pt"), map_location="cpu", weights_only=True) |
| self.meta = json.load(open(os.path.join(path, "meta.json"))) |
| self.model = Bind4(self.meta["dims"], self.meta["adapter"], self.meta.get("n_adapters", 1)) |
| self.model.load_state_dict(ck, strict=True) |
| self.model.eval().to(device) |
| self.hvg = np.asarray(self.meta["hvg_genes"], dtype=str) |
| self._enc = {} |
|
|
| |
| @torch.no_grad() |
| def he_from_feats(self, uni2, medvtok=None): |
| t = lambda a: None if a is None else torch.as_tensor(np.asarray(a, np.float32)).to(self.dev) |
| return nz(self.model.he(t(uni2), t(medvtok))).cpu() |
|
|
| @torch.no_grad() |
| def mif_from_feats(self, virtues): |
| return nz(self.model.mif(torch.as_tensor(np.asarray(virtues, np.float32)).to(self.dev))).cpu() |
|
|
| @torch.no_grad() |
| def rna_from_feats(self, feats, zscore="train"): |
| X = {e: torch.as_tensor(np.asarray(x, np.float32)).to(self.dev) for e, x in feats.items()} |
| return self.model.fuse("rna", X, zscore)[0].cpu() |
|
|
| @torch.no_grad() |
| def text_from_feats(self, feats): |
| X = {e: torch.as_tensor(np.asarray(x, np.float32)).to(self.dev) for e, x in feats.items()} |
| return self.model.fuse("text", X)[0].cpu() |
|
|
| |
| def _get(self, name): |
| if name in self._enc: |
| return self._enc[name] |
| if name == "uni2": |
| import timm |
| m = timm.create_model( |
| "vit_giant_patch14_224", pretrained=False, img_size=224, patch_size=14, depth=24, num_heads=24, |
| init_values=1e-5, embed_dim=1536, mlp_ratio=2.66667 * 2, num_classes=0, no_embed_class=True, |
| mlp_layer=timm.layers.SwiGLUPacked, act_layer=torch.nn.SiLU, reg_tokens=8, dynamic_img_size=True) |
| m.load_state_dict(torch.load(uni2_weights(), map_location="cpu", weights_only=True), strict=True) |
| obj = m.eval().to(self.dev) |
| elif name == "medvtok": |
| import json |
| sys.path.insert(0, os.path.join(ROOT, "ext", "medvtok", "repo")) |
| from models.medvtok import MedVTok |
| w = os.path.join(ROOT, "ext", "medvtok", "weights") |
| cfg = json.load(open(os.path.join(w, "config.json"))) |
| m = MedVTok(**cfg["model_kwargs"]) |
| m.load_state_dict(torch.load(os.path.join(w, "ckpt.pth"), map_location="cpu", weights_only=True), strict=True) |
| del m.decoder |
| obj = m.eval().to(self.dev) |
| elif name == "keep": |
| obj = _load_script("45b_text_feats_pathvlm.py").KEEPText().to(self.dev).eval() |
| elif name in ("biomedclip", "qwen3-0.6b"): |
| obj = _load_script("45_text_feats.py", [sys.argv[0]]).encoder(name)[0] |
| elif name == "tahoex1": |
| T = _load_script("43_rnafm_tahoex1.py") |
| obj = (T, *T.load_model()) |
| elif name == "se600m": |
| S = _load_script("43_rnafm_se600m.py") |
| obj = (S, *S.build()) |
| elif name == "aidocell": |
| A = _load_script("43_rnafm_aidocell.py") |
| obj = (A, *A.load_model()) |
| elif name == "virtues": |
| V = _load_script("14_extract_if_fm.py", [sys.argv[0], "virtues", "--no_probe"]) |
| uni = dict(V.MARKERS) |
| model, midx = V.load_virtues([uni[m] for m in MIF_MARKERS]) |
| obj = (V, model, midx.to(self.dev)) |
| else: |
| raise KeyError(name) |
| self._enc[name] = obj |
| return obj |
|
|
| def release(self): |
| """free the frozen encoders (GPU memory)""" |
| self._enc.clear(); torch.cuda.empty_cache() |
|
|
| |
| @torch.no_grad() |
| def he_feats(self, images, bs=32): |
| """UNI2-h CLS (224 px PIL bicubic, ImageNet norm, fp16) + MedVTok feats (256 px PIL bicubic, [-1, 1], bf16)""" |
| ims = _as_uint8_list(images) |
| uni, mv = self._get("uni2"), self._get("medvtok") |
| mean = torch.tensor([0.485, 0.456, 0.406], device=self.dev).view(1, 3, 1, 1) |
| std = torch.tensor([0.229, 0.224, 0.225], device=self.dev).view(1, 3, 1, 1) |
| U, M = [], [] |
| for s in range(0, len(ims), bs): |
| b = ims[s:s + bs] |
| x = torch.from_numpy(np.stack([np.asarray(i.resize((224, 224), 3)) for i in b])).to(self.dev) |
| x = (x.permute(0, 3, 1, 2).float() / 255 - mean) / std |
| with torch.autocast("cuda", dtype=torch.float16): |
| U.append(uni(x).float().cpu()) |
| x = torch.from_numpy(np.stack([np.asarray(i.resize((256, 256), 3)) for i in b])).to(self.dev) |
| x = x.permute(0, 3, 1, 2).float() / 255 * 2 - 1 |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| M.append(mv([i for i in x], sample_posterior=False, decode=False, return_batched_tensor=False)["feats"].float().cpu()) |
| return torch.cat(U), torch.cat(M) |
|
|
| def he(self, images, bs=32): |
| return self.he_from_feats(*self.he_feats(images, bs)) |
|
|
| @torch.no_grad() |
| def text_feats(self, texts, bs=64): |
| texts = ["" if t is None else str(t) for t in texts] |
| order = np.argsort([len(t) for t in texts]) |
| out = {} |
| for e in TXT_E: |
| f = self._get(e) |
| E = None |
| for i in range(0, len(texts), bs): |
| idx = order[i:i + bs] |
| if e == "keep": |
| z = f([texts[j] for j in idx]).float().cpu().numpy() |
| else: |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| z = f([texts[j] for j in idx]).float().cpu().numpy() |
| if E is None: |
| E = np.zeros((len(texts), z.shape[1]), np.float32) |
| E[idx] = z |
| out[e] = E.astype(np.float16).astype(np.float32) |
| return out |
|
|
| def text(self, texts, bs=64): |
| return self.text_from_feats(self.text_feats(texts, bs)) |
|
|
| def rna_feats(self, counts, genes, encoders=RNA_E): |
| import scipy.sparse as sp |
| X = sp.csr_matrix(counts, dtype=np.float32) |
| X.eliminate_zeros() |
| genes = np.asarray(genes, dtype=str) |
| out = {} |
| for e in encoders: |
| if e == "hvg": |
| tot = np.asarray(X.sum(1)).ravel() |
| L = (sp.diags(1e4 / np.clip(tot, 1, None)) @ X).tocsc() |
| L.data = np.log1p(L.data) |
| pos = {} |
| for j, g in enumerate(genes): |
| pos.setdefault(g, j) |
| cols = [pos.get(g, -1) for g in self.hvg] |
| H = np.zeros((X.shape[0], len(self.hvg)), np.float32) |
| ok = [i for i, c in enumerate(cols) if c >= 0] |
| H[:, ok] = L[:, [cols[i] for i in ok]].toarray() |
| out[e] = H |
| elif e == "tahoex1": |
| T, m, _, ccfg, vocab = self._get(e) |
| seqs, _ = T.make_seqs(X, genes, "user", vocab, ccfg, np.random.default_rng(0)) |
| out[e] = T.embed(m, seqs) |
| elif e == "se600m": |
| S, m, cfg, g2i = self._get(e) |
| Tk, C, _ = S.sentences(X, genes, g2i, cfg["dataset"]["pad_length"], cfg["dataset"]["cls_token_idx"], seed=0) |
| out[e] = S.embed(m, Tk, C) |
| elif e == "aidocell": |
| A, m, _ = self._get(e) |
| Xa, present = A.align(X, genes, "user") |
| out[e] = A.embed(m, Xa, present) |
| out[e] = np.asarray(out[e], np.float32).astype(np.float16).astype(np.float32) |
| return out |
|
|
| def rna(self, counts, genes, zscore="train", encoders=RNA_E): |
| return self.rna_from_feats(self.rna_feats(counts, genes, encoders), zscore) |
|
|
| @torch.no_grad() |
| def mif_feats(self, tiles, slide, size=168, bs=32): |
| """tiles (N, 16, H, W) in MIF_MARKERS order, 8-bit intensity scale (0-255); 14_ preprocessing: antialiased |
| bilinear resize to `size` px (~1 um/px for a 166 um tile), per-slide 0.99-quantile clip + log1p, mean / std |
| over the given tiles, 3x3 Gaussian blur; VirTues mean patch-summary token (fp16).""" |
| V, model, midx = self._get("virtues") |
| N = len(tiles) |
| X = np.empty((N, len(MIF_MARKERS), size, size), np.float16) |
| for s in range(0, N, 64): |
| x = torch.as_tensor(np.asarray(tiles[s:s + 64], np.float32)).to(self.dev) |
| X[s:s + 64] = F.interpolate(x, size=(size, size), mode="bilinear", antialias=True, |
| align_corners=False).clamp(min=0).half().cpu().numpy() |
| X = V.standardize(X, np.asarray(slide))[0] |
| out = np.zeros((N, 512), np.float32) |
| with torch.inference_mode(), torch.autocast("cuda", dtype=torch.float16): |
| for s in range(0, N, bs): |
| x = torch.from_numpy(X[s:s + bs]).to(self.dev).float() |
| ps = model.encoder.forward_list([c for c in x], [midx] * len(x)).patch_summary_tokens |
| out[s:s + bs] = torch.stack([p.float().mean((0, 1)) for p in ps]).cpu().numpy() |
| return out.astype(np.float16).astype(np.float32) |
|
|
| def mif(self, tiles, slide, size=168): |
| return self.mif_from_feats(self.mif_feats(tiles, slide, size)) |
|
|
