BigCell: Generating Gigapixel Whole-Slide Images

[📄 arXiv] [GitHub] [💾 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 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 and load the VAE from there. Requires diffusers>=0.36.

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.
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

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:

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

@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 foundation model and, for the CGM, the 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.
Downloads last month
-
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Collection including StonyBrook-CVLab/BigCell

Paper for StonyBrook-CVLab/BigCell