Blame for embed.py

"""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"]  # ORION panel without PD-1 (= 30_virtues_shared.py --m16)
UNI2_LOCAL = "/diskarray/home/tl688/reg2026/model_alldata/uni2-h.bin"  # override with $UNI2_PATH


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):  # = 42_ / 47_ / 48_ / 51_
    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():  # feature standardisation (identity where the encoder is used raw)
            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))  # True once statistics are set
        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)

    # ---- embeddings from (raw, un-normalised) encoder features; return un-normalised vectors
    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)  # PathGen PNGs: huge ICC profiles
    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 = {}

    # ------------------------------------------------------------------ feature-level
    @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()

    # ------------------------------------------------------------------ frozen encoders (lazy)
    def _get(self, name):
        if name in self._enc:
            return self._enc[name]
        if name == "uni2":  # = 04b_extract_anchor_fm.py
            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":  # = 50_medvtok_feats.py
            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()

    # ------------------------------------------------------------------ raw-input embedding
    @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])  # length-sorted batches as in 45_ / 45b_
        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":  # 45b_: fp32
                    z = f([texts[j] for j in idx]).float().cpu().numpy()
                else:  # 45_: bf16 autocast
                    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)  # the training features were stored as float16
        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":  # 21_hest_prep.py: log1p CP10k over the given genes, then the 2000 HEST HVGs (absent -> 0)
                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)  # stored as float16
        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))