PixCell: A Pan-Cancer Diffusion Foundation Model

         

---

We present PixCell, the first generative foundation model for digital histopathology. We progressively train our model to generate from 256x256 to 1024x1024 pixel images conditioned on UNI2-h embeddings. PixCell achieves state-of-the-art quality in digital pathology image generation and can be seamlessly used to perform targeted data augmentation and generative downstream tasks.

🔥 News

• Nov 2025: Released the virtual staining models + code.
• Sep 2025: Released training + sampling code for PixCell.
• June 2025: Released Hugging Face Diffusers checkpoints for PixCell:
  • PixCell-256 (Diffusers)
  • PixCell-1024 (Diffusers)

---

Contents

• 📜 Overview
• 🔧 Dependencies and Installation
• 📂 Dataset Preparation
• 🚀 Training
• 🔬 Sampling
• 🎛️ ControlNet
• 🎨 Virtual Staining
• 📦 Model Zoo
• 📄 Citation

🔧 Dependencies and Installation

• Python >= 3.9 (recommend Anaconda or Miniconda)
• PyTorch >= 2.0.1 + CUDA 11.7

conda create -n pixcell python=3.9
conda activate pixcell
conda install pytorch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 pytorch-cuda=11.7 -c pytorch -c nvidia

git clone https://github.com/<your-username>/pixcell.git
cd pixcell
pip install -r requirements.txt

📂 Dataset Preparation

Patch extraction

We train PixCell on ~70,000 Whole slide images (WSIs) from TCGA, CPTAC, GTeX, etc. WSIs should be preprocessed into patches using DSMIL.

Example structure:

tcga_diagnostic/
├── acc
│   └── single_1024
│       ├── TCGA-OR-A5J1-01Z-00-DX1
│       │   ├── 10_14.jpeg
│       │   ├── 10_15.jpeg
│       │   ├── 10_16.jpeg
│       │   └── 10_17.jpeg
...
cptac/
gtex/

To support progressive training, we only extract patches at 1024x1024 resolution. Lower resolutions are obtained by extracting crops from the 1024x1024 patches during training. See diffusion/data/datasets/pan_cancer.py for details.

Metadata

We index the datasets in an HDF5 file:

patches/metadata/patch_names_all.hdf5

import h5py
f = h5py.File('patches/metadata/patch_names_all.hdf5', 'r')

print(f.keys())
>>> <KeysViewHDF5 ['cptac_1024', 'gtex_1024', 'others_1024', 'sbu_olympus_1024', 'tcga_diagnostic_1024', 'tcga_fresh_frozen_1024']>

print(f['tcga_diagnostic_1024'][:5])
>>> array([b'tcga_diagnostic/lihc/single_1024/TCGA-BC-A8YO-01Z-00-DX1/82_17.jpeg',
       b'tcga_diagnostic/lihc/single_1024/TCGA-BC-A8YO-01Z-00-DX1/26_5.jpeg',
       b'tcga_diagnostic/lihc/single_1024/TCGA-BC-A8YO-01Z-00-DX1/64_11.jpeg',
       b'tcga_diagnostic/lihc/single_1024/TCGA-BC-A8YO-01Z-00-DX1/27_18.jpeg'],
      dtype=object)

Feature Extraction

For each patch, we pre-extract:

• VAE features (SD-3.5 VAE latents) — we use the VAE from stabilityai/stable-diffusion-3.5-large (the vae subfolder). Download it to pretrained_models/sd-3.5-vae.
• SSL embeddings (UNI2-h) — download the weights to pretrained_models/uni2/pytorch_model.bin.

Images live under patches/ and the extracted features live under a parallel features/ tree (mirroring the same relative paths):

<root>/
├── patches/
│   ├── metadata/patch_names_all.hdf5
│   └── <dataset>/<folder>/<patch>.jpeg
└── features/
    └── <dataset>/<folder>/
        ├── <patch>_sd3_vae.npy
        └── <patch>_uni.npy

Extract features with:

python tools/extract_features.py --root /path/to/dataset --dataset_name tcga_diagnostic --size 256

Sample dataset

To try sampling without preparing the full dataset, we release a small bundle of pre-extracted features (32 TCGA patches, in the layout above) on Hugging Face: StonyBrook-CVLab/PixCell-sample-data.

Download it and point data["root"] in your inference config.py at the extracted folder.

🚀 Training

1. Select a config file in configs/ (examples in configs/pan_cancer/).
   Example: configs/pan_cancer/pixart_20x_256.py

2. Launch training:

    accelerate launch train_scripts/train_pixcell.py configs/pan_cancer/pixart_20x_256.py \
        --work-dir /path/to/output_dir

Options:

• --work-dir : output directory for logs + checkpoints
• --resume-from : resume from a checkpoint
• --batch-size : override batch size

3. Configure accelerate (for multi-GPU/multi-node):

    accelerate config

🔄 Progressive Training

PixCell is trained in a progressive fashion to improve stability and efficiency (Similar to PixArt-Sigma):

1. Stage 1 — Train PixCell-256:
   Train the base model on 256×256 patches.

2. Stage 2 — Fine-tune PixCell-512:
   Initialize from PixCell-256 and fine-tune on 512×512 patches.

3. Stage 3 — Fine-tune PixCell-1024:
   Initialize from PixCell-512 and fine-tune on 1024×1024 patches.

Each stage reuses the weights of the previous resolution, allowing faster convergence.

See the config files in configs/pan_cancer/ for details.

🔬 Sampling

We provide sampling scripts for generating 256×256 (PixCell-256) and 1024×1024 (PixCell-1024) patches.

Setting up the --workdir. The scripts read a config.py from --workdir to build the model. For the released checkpoints, copy the matching ready-to-use config into your workdir and edit the two paths inside it (vae_pretrained, data["root"]):

mkdir -p /path/to/workdir/checkpoints
# PixCell-256:
cp configs/pan_cancer/pixcell_256_inference.py  /path/to/workdir/config.py
# PixCell-1024:
cp configs/pan_cancer/pixcell_1024_inference.py /path/to/workdir/config.py

Then place the downloaded checkpoint at the path you pass to --checkpoint. The released pixcell_256.ckpt / pixcell_1024.ckpt are the EMA weights (stored under the state_dict key); the scripts just default to checkpoints/last_ema.ckpt, so point --checkpoint at your downloaded file (e.g. checkpoints/pixcell_256.ckpt).

The UNI conditioning embeddings come from a dataset of pre-extracted features. To try sampling without the full dataset, use the small sample dataset and set data["root"] in your config.py to it.

Example for 256×256 generation:

python tools/sample_256.py \
    --workdir /path/to/workdir \
    --checkpoint checkpoints/last_ema.ckpt \
    --out_dir samples_256 \
    --n_images 5000 \
    --sampling_steps 20 \
    --guidance_strength 2 \
    --sampling_algo dpm-solver

Outputs:

samples_256/
├── real/   # real patches (for FID evaluation)
├── syn/    # generated synthetic patches

Using our CVPR24 Large image generation algorithm, we can generate 4K×4K images with PixCell-1024:

python tools/sample_4k.py \
    --workdir /path/to/workdir \
    --uni_weights pretrained_models/uni2/pytorch_model.bin \
    --image_list /path/to/image_list.npy \
    --output_dir samples_4k \
    --num_samples 100 \
    --num_timesteps 20 \
    --guidance_scale 2 \
    --sliding_window_size 64 \
    --gpu_id 0

4K generation is reference-based: image_list.npy is a NumPy array of string paths to real 4096×4096 images, each of which is diced into 256×256 patches to build the UNI conditioning grid.

Hugging Face Diffusers Sampling

We also provide Diffusers-compatible checkpoints and sampling code on Hugging Face:

• PixCell-256 (Diffusers)
• PixCell-1024 (Diffusers)

Follow the instructions on those pages to sample using the diffusers API.

🎛️ ControlNet Training

We provide a mock implementation of our ControlNet training pipeline in controlnet/train.py.

You will need to replace the dataloader with your own image condition dataset. Note that the ControlNet transformer also requires a UNI-2h embedding as conditioning in our implementation. It may also be possible to train without any conditioning.

🎨 Virtual Staining

Inference

We provide the implementation of our virtual staining algorithm in a Jupyter notebook virtual_staining.ipynb.

The virtual staining relies on additional model weights (PixCell-1024 LoRA, flow-matching MLP). For the four stains of the MIST dataset (HER2, ER, PR, Ki67) and the HER2Match dataset, the notebook downloads the necessary models from our Huggingface repository.

Training

We provide scripts for training the PixCell-1024 IHC LoRA and the flow-matching MLP.

For LoRA training, you can run the script using accelerate:

CUDA_VISIBLE_DEVICES=1 accelerate launch --num_processes N \
    virtual_staining/train_lora.py \
    --dataset [MIST/HER2Match] \
    --root_dir /path/to/data/ \
    --split train \
    --stain [HER2/PR/ER/Ki67/ ] \
    --train_batch_size 4 \
    --num_epochs 10 \
    --gradient_accumulation_steps 2

For the MLP training, you can run it as:

python virutal_staining/train_flow_mlp.py \
    --dataset [MIST/HER2Match] \
    --root_dir /path/to/data/ \
    --split train \
    --stain [HER2/PR/ER/Ki67/ ] \
    --device cuda \
    --train_batch_size 4 \
    --num_epochs 100 \
    --save_every 25 \

To speed up training for both the LoRA and MLP, we suggest pre-extracting the UNI embeddings for the images in the dataset. The scripts we provide assume that the UNI embeddings are not pre-extracted.

📦 Model Zoo

---

📄 Citation

If you use PixCell in your research, please cite:

@misc{yellapragada2025pixcell,
    title={PixCell: A generative foundation model for digital histopathology images}, 
    author={Srikar Yellapragada and Alexandros Graikos and Zilinghan Li and Kostas Triaridis and Varun Belagali and Tarak Nath Nandi and Karen Bai and Beatrice S. Knudsen and Tahsin Kurc and Rajarsi R. Gupta and Prateek Prasanna and Ravi K Madduri and Joel Saltz and Dimitris Samaras},
    year={2025},
    eprint={2506.05127},
    archivePrefix={arXiv},
    primaryClass={eess.IV},
    url={https://arxiv.org/abs/2506.05127}, 
}

---

🤗 Acknowledgements

PixCell builds on PixArt-Sigma and Diffusers.