| --- |
| license: mit |
| library_name: pytorch |
| tags: |
| - pathology |
| - computational-pathology |
| - virtual-staining |
| - multiplex-immunofluorescence |
| - diffusion |
| - flow-matching |
| - image-to-image |
| base_model: xiangjx/MuPaD-512 |
| pipeline_tag: image-to-image |
| --- |
| |
| # MuPaD-HE2mIF |
|
|
| Predict a **16-channel Orion multiplexed immunofluorescence (mIF) panel from a |
| single H&E tile**, by finetuning [MuPaD-512](https://huggingface.co/xiangjx/MuPaD-512) |
| as a conditional **mean** predictor. |
|
|
| Code: [github.com/jinxixiang/MUPAD](https://github.com/jinxixiang/MUPAD) → `5-HE2mIF/` |
|
|
| ## Why a mean predictor |
|
|
| The evaluation metric is a per-channel Pearson correlation against ground truth, |
| and that is maximized by the conditional mean `E[mIF | H&E]` — not by a sharp |
| sample from `p(mIF | H&E)`. The model is trained with flow matching but **read |
| out at `t = 1`**, where the input is pure noise and the loss-optimal output *is* |
| the conditional mean. Inference is a single forward pass, not a multi-step |
| sampler. |
|
|
| ``` |
| z_he = LatentNorm_he( HE_VAE(he) ) # pretrained H&E latent, frozen |
| z_mif = LatentNorm_mif( mIF_VAE(mif) ) # stage-1 mIF latent, frozen |
| z_t = (1-t)·z_mif + t·ε |
| x̂ = SiT( [z_t ; z_he], t, MUSK(he) ) # x-prediction |
| |
| inference: x̂ = SiT( [ε ; z_he], t=1, MUSK(he) ) → decode |
| ``` |
|
|
| H&E enters as **input, not context**: `z_he` is concatenated channel-wise at the |
| patch embedder, and the new condition slots are initialized by copying the |
| pretrained patch-embed weights, so `z_he` lives in exactly the latent space the |
| trunk was pretrained on. MUSK cross-attention is a secondary semantic channel. |
|
|
| ## Contents |
|
|
| ``` |
| stage1_vae/ |
| vae/config.json |
| vae/diffusion_pytorch_model.safetensors 16-channel mIF VAE (335 MB) |
| latent_stats.json mIF latent mean/var, fitted at 256px |
| latent_stats_512.json same, fitted at 512px <- used at inference |
| mif_input_scale.json raw-intensity scaling used by the VAE |
| config.yaml stage-1 training config |
| stage2_denoising/ |
| diffusion_pytorch_model.safetensors SiT trunk, 1.39B params, fp32 (5.56 GB) |
| config.yaml stage-2 training config |
| ``` |
|
|
| **Latent statistics are resolution-specific.** `latent_stats.json` was fitted at |
| the VAE's own 256px; applying it at 512px leaves normalized latents at std |
| 0.68–0.79 instead of 1.0, which unbalances the interpolant. Use |
| `latent_stats_512.json` for 512px inference — the code selects it automatically. |
|
|
| ## Requires |
|
|
| This repository holds only the HE2mIF-specific weights. You also need |
| [**xiangjx/MuPaD-512**](https://huggingface.co/xiangjx/MuPaD-512) for the SiT |
| trunk architecture, the H&E VAE, and the MUSK encoder. |
|
|
| ## Usage |
|
|
| ```bash |
| git clone https://github.com/jinxixiang/MUPAD && cd MUPAD/5-HE2mIF |
| pip install -r ../requirements.txt |
| |
| huggingface-cli download xiangjx/MuPaD-512 --local-dir pretrained/MuPaD-512 |
| huggingface-cli download xiangjx/MuPaD-HE2mIF --local-dir exps |
| export MUPAD_PRETRAINED_ROOT=$PWD/pretrained/MuPaD-512 |
| export ORION_DATA_ROOT=/path/to/ORIONCRC_dataset_tile_20x |
| ``` |
|
|
| ```bash |
| python generate.py \ |
| --config configs/stage2_denoising.yaml \ |
| --denoise-ckpt exps/stage2_denoising/diffusion_pytorch_model.safetensors \ |
| --vae-ckpt exps/stage1_vae \ |
| --val-csv $ORION_DATA_ROOT/val_dataframe.csv \ |
| --out-dir predictions/run \ |
| --guidance-scale 1.25 |
| |
| python eval/evaluate.py --pred-dir predictions/run \ |
| --csv $ORION_DATA_ROOT/val_dataframe.csv --out results.csv |
| ``` |
|
|
| Guidance uses zeroed conditioning as the null branch (there is no conditioning |
| dropout in training). `--guidance-scale 1.25` measurably beats the unguided |
| default. |
|
|
| ## Markers |
|
|
| 16 of Orion's 17 channels. `PD-1` is excluded: its positivity column is 1.0000 |
| on every training tile — a broken threshold, not biology — so it carries no |
| usable supervision. |
|
|
| `Hoechst`, `CD31`, `CD45`, `CD68`, `CD4`, `FOXP3`, `CD8a`, `CD45RO`, `CD20`, |
| `PD-L1`, `CD3e`, `CD163`, `E-cadherin`, `Ki67`, `Pan-CK`, `SMA` |
|
|
| ## Data |
|
|
| Trained on **ORION-CRC** tiles at 20× (512 × 512 px), |
| [Zenodo record 15340874](https://zenodo.org/records/15340874). |
|
|
| ## License |
|
|
| MIT, following the REPA/SiT lineage the trunk derives from. |
|
|
