README.md

---
license: cc-by-nc-nd-4.0
library_name: diffusers
tags:
- histopathology
- whole-slide-image
- flow-matching
---

# BigCell: Generating Gigapixel Whole-Slide Images

[[📄 arXiv]](https://arxiv.org/abs/2610.14261) [[GitHub]](https://github.com/cvlab-stonybrook/BigCell) [[💾 Sample data]](https://huggingface.co/datasets/StonyBrook-CVLab/BigCell-sample-data)

BigCell is a flow-matching transformer (SiT-XL/2, 845M parameters) that generates histopathology images at 20×.
It is conditioned on a dense grid of [UNI-2h](https://huggingface.co/MahmoodLab/UNI2-h) embeddings, with one
1536-d token per 256×256 px. The model works in the latent space of the FLUX.2 VAE.

A native 2048×2048 tile generator, used with an inference-time sliding window, produces coherent images of up to
32,768×32,768 px. A separate conditioning generation model (CGM) samples the UNI-2 grid from a pathology report,
for text-to-image generation at 8192×8192.

| Folder | Content | Native tile | UNI-2 tokens / tile |
|---|---|---|---|
| `bigcell-256` | BigCell tile generator, progressive-training stage 1 | 256 px | 1 |
| `bigcell-1024` | stage 3 | 1024 px | 4×4 |
| `bigcell-2048` | **final model** (used in the paper) | 2048 px | 8×8 |
| `cgm-8192` | CGM (DiT-DH-B): TITAN text → 32×32×1536 UNI-2 grid of an 8192 px region | – | – |

Each `bigcell-*` folder is a diffusers pipeline (`BigCellPipeline`).

## Usage

The FLUX.2 VAE weights are not hosted here. Accept the licence of
[black-forest-labs/FLUX.2-dev](https://huggingface.co/black-forest-labs/FLUX.2-dev) and load the VAE from there.
Requires `diffusers>=0.36`.

```python
import torch
from diffusers import AutoencoderKLFlux2, DiffusionPipeline
from huggingface_hub import snapshot_download

name = "bigcell-2048"
path = f"{snapshot_download('StonyBrook-CVLab/BigCell', allow_patterns=f'{name}/*')}/{name}"

vae = AutoencoderKLFlux2.from_pretrained("black-forest-labs/FLUX.2-dev", subfolder="vae", torch_dtype=torch.bfloat16)
pipe = DiffusionPipeline.from_pretrained(path, vae=vae, custom_pipeline=path, trust_remote_code=True,
                                         torch_dtype=torch.bfloat16).to("cuda")
```

### Generate from a UNI-2 grid

The pipeline takes a UNI-2 grid of shape `(B, Hg, Wg, 1536)` and returns a `(256·Hg) × (256·Wg)` image.

- A grid of exactly one native tile (8×8 for `bigcell-2048`) is denoised directly. Tokens of shape
  `(B, T, 1536)` in raster order are also accepted.
- Larger grids are generated with the sliding-window procedure of the paper. Set `top_r < 1` for adaptive window
  scheduling: only the fraction `top_r` of windows with the largest classifier-free-guidance magnitude is
  recomputed after the first step.

```python
import numpy as np
from huggingface_hub import hf_hub_download

# 32×32 UNI-2 grid of a real 8192×8192 region from the sample data
f = hf_hub_download("StonyBrook-CVLab/BigCell-sample-data", repo_type="dataset",
                    filename="cptac/AML/C3L-00453-41_x0_y0/8192/0_uni.npy")
grid = torch.from_numpy(np.load(f)).float()[None]            # (1, 32, 32, 1536)

image = pipe(grid, num_inference_steps=20, guidance_scale=2.0, top_r=0.25,
             generator=torch.Generator("cuda").manual_seed(0)).images[0]   # 8192×8192 PIL image
```

`pipe.get_unconditional_embedding(B)` returns the learned null condition, for unconditional samples.

### Extract a UNI-2 grid from your own image

```python
import timm, einops
from PIL import Image
from timm.data import resolve_data_config
from timm.data.transforms_factory import create_transform

timm_kwargs = dict(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)
uni = timm.create_model("hf-hub:MahmoodLab/UNI2-h", pretrained=True, **timm_kwargs).eval().cuda()
transform = create_transform(**resolve_data_config(uni.pretrained_cfg, model=uni))

img = np.array(Image.open("region_20x.png").convert("RGB"))        # H×W, multiples of 256, at 20×
hg, wg = img.shape[0] // 256, img.shape[1] // 256
patches = einops.rearrange(img, "(hg h) (wg w) c -> (hg wg) h w c", h=256, w=256)
with torch.inference_mode():
    feats = torch.cat([uni(torch.stack([transform(Image.fromarray(p)) for p in patches[i:i + 256]]).cuda())
                       for i in range(0, len(patches), 256)])
grid = feats.reshape(1, hg, wg, 1536).float().cpu()
```

### Multi-GPU 32k generation and text-to-image

Gigapixel generation with windows sharded across GPUs, text-to-image with the CGM, and CGM-based data
augmentation are in the [GitHub repository](https://github.com/cvlab-stonybrook/BigCell):

```bash
torchrun --nproc_per_node 4 scripts/generate_from_grid.py --grid .../32768/0_uni.npy --top-r 0.25 --out 32k.jpg
python scripts/generate_from_text.py --prompt "Pancreatic ductal adenocarcinoma ..." --out out/
```

## Sampling settings used in the paper

- BigCell: 20 Euler steps, CFG 2.0, 2048 px windows with a 1024 px stride, and Gaussian window weighting.
- The time shift is stored in `transformer/config.json` (`alpha_shift`) and `scheduler/scheduler_config.json`
  (`shift`). It is 3.0 for `bigcell-1024` / `bigcell-2048` and 1.0 for `bigcell-256`.
- CGM: 50-point Euler schedule with time shift √(32·32·1536/4096), CFG 2.0, and TITAN text tokens (unpooled,
  127×768) as the condition.

## Citation

```bibtex
@misc{yellapragada2026bigcell,
      title={BigCell: Generating Gigapixel Whole-Slide Images}, 
      author={Srikar Yellapragada and Alexandros Graikos and Zilinghan Li and Kostas Triaridis and Tarak Nath Nandi and Ji Dong K Bai and Beatrice Knudsen and Tahsin Kurc and Prateek Prasanna and Rajarsi R. Gupta and Ravi K Madduri and Joel Saltz and Dimitris Samaras},
      year={2026},
      eprint={2610.14261},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2610.14261}, 
}
```

## License & Usage

**License**: Creative Commons Attribution-NonCommercial-NoDerivatives 4.0 International (CC BY-NC-ND 4.0)

**Notice**: These models are derivative works conditioned on embeddings from the
[UNI-2h](https://huggingface.co/MahmoodLab/UNI2-h) foundation model and, for the CGM, the
[TITAN](https://huggingface.co/MahmoodLab/TITAN) text encoder. They are subject to the original terms of those
licences. The FLUX.2 VAE is governed by its own licence.

- Academic and research use only: you may use these weights for non-commercial research purposes.
- No commercial use, including product development or commercial services.
- Generated images must not be used for clinical diagnosis.