Add files using upload-large-folder tool

README.md

@@ -0,0 +1,141 @@
+---
+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.

bigcell-1024/model_index.json

@@ -0,0 +1,19 @@
+{
+  "_class_name": [
+    "pipeline",
+    "BigCellPipeline"
+  ],
+  "_diffusers_version": "0.37.1",
+  "transformer": [
+    "bigcell_transformer_2d",
+    "BigCellTransformer2DModel"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKLFlux2"
+  ],
+  "scheduler": [
+    "diffusers",
+    "FlowMatchEulerDiscreteScheduler"
+  ]
+}

bigcell-1024/pipeline.py

@@ -0,0 +1,151 @@
+# BigCell pipeline. https://github.com/cvlab-stonybrook/BigCell
+#
+# Generates histopathology images from a grid of UNI-2 embeddings (one 1536-d token per 256x256 px at 20x).
+# A grid that matches the model's native tile (e.g. 8x8 for BigCell-2048) is denoised directly; larger grids
+# are generated with the sliding-window procedure of the paper (BigCell-large), optionally with adaptive
+# window scheduling (``top_r < 1``). For multi-GPU generation of 16k-32k images use the GitHub scripts.
+
+import math
+from typing import List, Optional, Union
+
+import numpy as np
+import torch
+from diffusers import DiffusionPipeline
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.pipelines.pipeline_utils import ImagePipelineOutput
+from diffusers.utils.torch_utils import randn_tensor
+from PIL import Image
+
+LATENTS_PER_TOKEN = 32  # one UNI-2 token covers 256 px = 32 Flux.2 latents
+
+
+def _gaussian_kernel(size, sigma, device):
+    x = torch.linspace(-1, 1, size)
+    x = torch.stack((x.tile(size, 1), x.tile(size, 1).T), dim=0)
+    k = torch.exp(-(torch.linalg.norm(x, dim=0) ** 2) / (2 * sigma**2))
+    return (k / k.max()).to(device)
+
+
+class BigCellPipeline(DiffusionPipeline):
+    model_cpu_offload_seq = "transformer->vae"
+
+    def __init__(self, transformer, vae, scheduler):
+        super().__init__()
+        self.register_modules(transformer=transformer, vae=vae, scheduler=scheduler)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=8)
+
+    @property
+    def tile_tokens(self):
+        """Side of the UNI-2 grid covered by one native tile (1, 4 or 8)."""
+        return int(round(math.sqrt(self.transformer.config.num_caption_tokens)))
+
+    def get_unconditional_embedding(self, batch_size=1):
+        return self.transformer.unconditional_embedding.tile(batch_size, 1, 1)
+
+    def _velocity(self, x, t, y, guidance_scale):
+        """Returns (cfg velocity, importance score |v_cond - v_uncond|)."""
+        null = self.get_unconditional_embedding(x.shape[0]).to(y.dtype)
+        v = self.transformer(torch.cat([x, x]), t.expand(2 * x.shape[0]), torch.cat([y, null])).float()
+        v_cond, v_uncond = v.chunk(2)
+        return v_uncond + guidance_scale * (v_cond - v_uncond), (v_cond - v_uncond).abs().mean().item()
+
+    def _sigmas(self, num_inference_steps, device):
+        self.scheduler.set_timesteps(sigmas=np.linspace(1.0, 0.0, num_inference_steps + 1)[:-1], device=device)
+        return self.scheduler.sigmas.to(device)  # shifted, ends with 0
+
+    @torch.no_grad()
+    def _decode(self, latents, tile=256, sigma=0.5):
+        """Tiled decode for large canvases (Gaussian-blended, accumulated on CPU)."""
+        _, _, h, w = latents.shape
+        dtype = self.vae.dtype
+        if h <= tile and w <= tile:
+            return self.vae.decode(latents.to(dtype)).sample.float().cpu()
+        out = torch.zeros(latents.shape[0], 3, h * 8, w * 8)
+        cnt = torch.zeros(1, 1, h * 8, w * 8)
+        kernel = _gaussian_kernel(tile * 8, sigma, "cpu")[None, None]
+        stride = tile // 2
+        for i in range(0, h - tile + 1, stride):
+            for j in range(0, w - tile + 1, stride):
+                dec = self.vae.decode(latents[:, :, i:i + tile, j:j + tile].to(dtype)).sample.float().cpu()
+                out[:, :, i * 8:(i + tile) * 8, j * 8:(j + tile) * 8] += dec * kernel
+                cnt[:, :, i * 8:(i + tile) * 8, j * 8:(j + tile) * 8] += kernel
+        return out / cnt.clamp(min=1e-6)
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        uni_embeds: torch.Tensor,
+        num_inference_steps: int = 20,
+        guidance_scale: float = 2.0,
+        top_r: float = 1.0,
+        stride: Optional[int] = None,
+        window_sigma: float = 0.5,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+        progress: bool = True,
+    ):
+        """
+        Args:
+            uni_embeds: UNI-2 grid of shape (B, Hg, Wg, 1536), or tokens of one native tile (B, T, 1536).
+                The output image is (256 * Hg) x (256 * Wg) pixels.
+            num_inference_steps: Euler steps (paper: 20).
+            guidance_scale: classifier-free guidance scale (paper: 2.0).
+            top_r: fraction of sliding windows recomputed at every step after the first (paper Sec. 3.3);
+                the rest reuse their cached velocity. 1.0 = full computation.
+            stride: window stride in latents (default: half a window, e.g. 128 = 1024 px for BigCell-2048).
+            window_sigma: width of the Gaussian window weighting.
+            output_type: "pil", "np" (uint8 HxWx3) or "latent".
+        """
+        device = self._execution_device
+        dtype = self.transformer.dtype
+        tok = self.tile_tokens
+        if uni_embeds.ndim == 3:  # (B, T, C) tokens of one tile -> grid
+            uni_embeds = uni_embeds.reshape(uni_embeds.shape[0], tok, tok, -1)
+        B, Hg, Wg, _ = uni_embeds.shape
+        if Hg < tok or Wg < tok:
+            raise ValueError(f"grid {Hg}x{Wg} is smaller than one tile ({tok}x{tok})")
+        uni_embeds = uni_embeds.to(device, dtype)
+
+        win = self.transformer.config.sample_size
+        H, W = Hg * LATENTS_PER_TOKEN, Wg * LATENTS_PER_TOKEN
+        stride = stride or win // 2
+        windows = [(h, w) for h in range(0, H - win + 1, stride) for w in range(0, W - win + 1, stride)]
+        kernel = _gaussian_kernel(win, window_sigma, device)
+
+        if latents is None:
+            latents = randn_tensor((B, self.transformer.config.in_channels, H, W), generator=generator,
+                                   device=device, dtype=torch.float32)
+        x = latents.to(device, torch.float32)
+        sigmas = self._sigmas(num_inference_steps, device)
+
+        n_select = len(windows) if top_r >= 1 else max(1, math.ceil(top_r * len(windows)))
+        importance = [dict() for _ in range(B)]
+        v_prev = torch.zeros_like(x)
+        with self.progress_bar(total=num_inference_steps) as pbar:
+            for step, (s_cur, s_next) in enumerate(zip(sigmas[:-1], sigmas[1:])):
+                t = s_cur.reshape(1).to(dtype)
+                v_new = v_prev.clone()
+                for b in range(B):
+                    active = windows if step == 0 else sorted(windows, key=lambda w: -importance[b].get(w, 0.0))[:n_select]
+                    v_sum = torch.zeros_like(x[b:b + 1])
+                    w_sum = torch.zeros_like(x[b:b + 1])
+                    for h, w in active:
+                        y = uni_embeds[b:b + 1, h // 32:h // 32 + tok, w // 32:w // 32 + tok].reshape(1, tok * tok, -1)
+                        v, score = self._velocity(x[b:b + 1, :, h:h + win, w:w + win].to(dtype), t, y, guidance_scale)
+                        importance[b][(h, w)] = score
+                        v_sum[:, :, h:h + win, w:w + win] += v * kernel
+                        w_sum[:, :, h:h + win, w:w + win] += kernel
+                    v_new[b:b + 1] = torch.where(w_sum > 0, v_sum / w_sum.clamp(min=1e-6), v_prev[b:b + 1])
+                v_prev = v_new
+                x = x + (s_next - s_cur) * v_new
+                if progress:
+                    pbar.update()
+
+        if output_type == "latent":
+            return ImagePipelineOutput(images=x) if return_dict else (x,)
+        img = ((self._decode(x) + 1) / 2).clamp(0, 1)
+        img = (img.permute(0, 2, 3, 1) * 255).round().to(torch.uint8).numpy()
+        images = [Image.fromarray(i) for i in img] if output_type == "pil" else img
+        return ImagePipelineOutput(images=images) if return_dict else (images,)

bigcell-1024/scheduler/scheduler_config.json

@@ -0,0 +1,18 @@
+{
+  "_class_name": "FlowMatchEulerDiscreteScheduler",
+  "_diffusers_version": "0.37.1",
+  "base_image_seq_len": 256,
+  "base_shift": 0.5,
+  "invert_sigmas": false,
+  "max_image_seq_len": 4096,
+  "max_shift": 1.15,
+  "num_train_timesteps": 1000,
+  "shift": 3.0,
+  "shift_terminal": null,
+  "stochastic_sampling": false,
+  "time_shift_type": "exponential",
+  "use_beta_sigmas": false,
+  "use_dynamic_shifting": false,
+  "use_exponential_sigmas": false,
+  "use_karras_sigmas": false
+}

bigcell-1024/transformer/bigcell_transformer_2d.py

@@ -0,0 +1,202 @@
+# BigCell tile generator: a SiT flow-matching transformer conditioned on a grid of UNI-2 embeddings
+# through cross-attention. https://github.com/cvlab-stonybrook/BigCell
+#
+# Adapted from SiT (https://github.com/willisma/SiT, MIT) and REPA-E (https://github.com/End2End-Diffusion/REPA-E, MIT).
+
+import math
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+class PatchEmbed(nn.Module):
+    def __init__(self, patch_size, in_channels, hidden_size):
+        super().__init__()
+        self.proj = nn.Conv2d(in_channels, hidden_size, kernel_size=patch_size, stride=patch_size, bias=True)
+
+    def forward(self, x):
+        return self.proj(x).flatten(2).transpose(1, 2)
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features, out_features=None, approximate="none"):
+        super().__init__()
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = nn.GELU(approximate=approximate)
+        self.fc2 = nn.Linear(hidden_features, out_features or in_features)
+
+    def forward(self, x):
+        return self.fc2(self.act(self.fc1(x)))
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads):
+        super().__init__()
+        self.num_heads = num_heads
+        self.qkv = nn.Linear(dim, dim * 3, bias=True)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        q, k, v = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4).unbind(0)
+        x = F.scaled_dot_product_attention(q, k, v)
+        return self.proj(x.transpose(1, 2).reshape(B, N, C))
+
+
+class CrossAttention(nn.Module):
+    def __init__(self, dim, num_heads):
+        super().__init__()
+        self.num_heads = num_heads
+        self.q_linear = nn.Linear(dim, dim)
+        self.kv_linear = nn.Linear(dim, dim * 2)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x, cond):
+        B, N, C = x.shape
+        h = self.num_heads
+        q = self.q_linear(x).reshape(B, N, h, C // h).transpose(1, 2)
+        k, v = self.kv_linear(cond).reshape(B, cond.shape[1], 2, h, C // h).permute(2, 0, 3, 1, 4).unbind(0)
+        x = F.scaled_dot_product_attention(q, k, v)
+        return self.proj(x.transpose(1, 2).reshape(B, N, C))
+
+
+class TimestepEmbedder(nn.Module):
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    def forward(self, t):
+        half = self.frequency_embedding_size // 2
+        freqs = torch.exp(-math.log(10000) * torch.arange(half, dtype=torch.float32, device=t.device) / half)
+        args = t[:, None].float() * freqs[None]
+        return self.mlp(torch.cat([torch.cos(args), torch.sin(args)], dim=-1).to(t.dtype))
+
+
+class CaptionEmbedder(nn.Module):
+    """Projects UNI-2 tokens; ``y_embedding`` is the learned unconditional embedding used for CFG."""
+
+    def __init__(self, in_channels, hidden_size, token_num):
+        super().__init__()
+        self.y_proj = Mlp(in_channels, hidden_size, hidden_size, approximate="tanh")
+        self.register_buffer("y_embedding", torch.randn(token_num, in_channels) / in_channels**0.5)
+
+    def forward(self, caption):
+        return self.y_proj(caption)
+
+
+class BigCellBlock(nn.Module):
+    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.attn = Attention(hidden_size, num_heads)
+        self.cross_attn = CrossAttention(hidden_size, num_heads)
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.mlp = Mlp(hidden_size, int(hidden_size * mlp_ratio), approximate="tanh")
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True))
+
+    def forward(self, x, cond, t):
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(t).chunk(6, dim=-1)
+        x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa))
+        x = x + self.cross_attn(x, cond)
+        x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
+        return x
+
+
+class FinalLayer(nn.Module):
+    def __init__(self, hidden_size, patch_size, out_channels):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
+
+    def forward(self, x, c):
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)
+        return self.linear(modulate(self.norm_final(x), shift, scale))
+
+
+class BigCellTransformer2DModel(ModelMixin, ConfigMixin):
+    """SiT-XL/2 velocity predictor operating on Flux.2 VAE latents (32 channels, 8x downsampling).
+
+    Args:
+        sample_size: latent side length (tile resolution // 8).
+        num_caption_tokens: number of UNI-2 tokens per tile, (tile resolution // 256) ** 2.
+        alpha_shift: time shift used during training; the scheduler uses the same value.
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        sample_size: int = 256,
+        patch_size: int = 2,
+        in_channels: int = 32,
+        hidden_size: int = 1152,
+        num_layers: int = 28,
+        num_attention_heads: int = 16,
+        mlp_ratio: float = 4.0,
+        caption_channels: int = 1536,
+        num_caption_tokens: int = 64,
+        alpha_shift: float = 1.0,
+    ):
+        super().__init__()
+        self.x_embedder = PatchEmbed(patch_size, in_channels, hidden_size)
+        self.t_embedder = TimestepEmbedder(hidden_size)
+        self.y_embedder = CaptionEmbedder(caption_channels, hidden_size, num_caption_tokens)
+        num_patches = (sample_size // patch_size) ** 2
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, hidden_size), requires_grad=False)
+        grid = int(num_caption_tokens**0.5)
+        if num_caption_tokens > 1 and grid**2 == num_caption_tokens:
+            self.y_pos_embed = nn.Parameter(torch.zeros(1, num_caption_tokens, hidden_size), requires_grad=False)
+        else:
+            self.y_pos_embed = None
+        self.blocks = nn.ModuleList(
+            [BigCellBlock(hidden_size, num_attention_heads, mlp_ratio) for _ in range(num_layers)]
+        )
+        self.final_layer = FinalLayer(hidden_size, patch_size, in_channels)
+        pos = _sincos_2d(hidden_size, np.stack(np.meshgrid(np.arange(sample_size // patch_size),
+                                                          np.arange(sample_size // patch_size)), 0).reshape(2, -1))
+        self.pos_embed.data.copy_(torch.from_numpy(pos).float()[None])
+
+    @property
+    def unconditional_embedding(self):
+        """(1, T, 1536) learned null condition for classifier-free guidance."""
+        return self.y_embedder.y_embedding[None]
+
+    def unpatchify(self, x):
+        c, p = self.config.in_channels, self.config.patch_size
+        h = w = int(x.shape[1] ** 0.5)
+        x = torch.einsum("nhwpqc->nchpwq", x.reshape(x.shape[0], h, w, p, p, c))
+        return x.reshape(x.shape[0], c, h * p, w * p)
+
+    def forward(self, hidden_states, timestep, encoder_hidden_states):
+        """hidden_states: (B, 32, H, W) noisy latent; timestep: (B,) in [0, 1] (1 = pure noise);
+        encoder_hidden_states: (B, T, 1536) UNI-2 tokens in raster order. Returns the velocity (noise - data)."""
+        x = self.x_embedder(hidden_states) + self.pos_embed
+        t = self.t_embedder(timestep)
+        y = self.y_embedder(encoder_hidden_states)
+        if self.y_pos_embed is not None:
+            y = y + self.y_pos_embed
+        for block in self.blocks:
+            x = block(x, y, t)
+        return self.unpatchify(self.final_layer(x, t))
+
+
+def _sincos_2d(embed_dim, grid):
+    def one_d(dim, pos):
+        omega = 1.0 / 10000 ** (np.arange(dim // 2, dtype=np.float64) / (dim / 2.0))
+        out = np.einsum("m,d->md", pos.reshape(-1), omega)
+        return np.concatenate([np.sin(out), np.cos(out)], axis=1)
+
+    return np.concatenate([one_d(embed_dim // 2, grid[0]), one_d(embed_dim // 2, grid[1])], axis=1)

bigcell-1024/transformer/config.json

@@ -0,0 +1,14 @@
+{
+  "_class_name": "BigCellTransformer2DModel",
+  "_diffusers_version": "0.37.1",
+  "alpha_shift": 3.0,
+  "caption_channels": 1536,
+  "hidden_size": 1152,
+  "in_channels": 32,
+  "mlp_ratio": 4.0,
+  "num_attention_heads": 16,
+  "num_caption_tokens": 16,
+  "num_layers": 28,
+  "patch_size": 2,
+  "sample_size": 128
+}

bigcell-1024/transformer/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7e13cd2c48e9a4e7673e1cc9da626ee9842e5abbb22ba7b6036993d5015567b1
+size 3322233744

bigcell-1024/vae/config.json

@@ -0,0 +1,40 @@
+{
+  "_class_name": "AutoencoderKLFlux2",
+  "_diffusers_version": "0.36.0.dev0",
+  "act_fn": "silu",
+  "batch_norm_eps": 0.0001,
+  "batch_norm_momentum": 0.1,
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "force_upcast": true,
+  "in_channels": 3,
+  "latent_channels": 32,
+  "layers_per_block": 2,
+  "mid_block_add_attention": true,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "patch_size": [
+    2,
+    2
+  ],
+  "sample_size": 1024,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ],
+  "use_post_quant_conv": true,
+  "use_quant_conv": true,
+  "_name_or_path": "black-forest-labs/FLUX.2-dev"
+}

bigcell-2048/model_index.json

@@ -0,0 +1,19 @@
+{
+  "_class_name": [
+    "pipeline",
+    "BigCellPipeline"
+  ],
+  "_diffusers_version": "0.37.1",
+  "transformer": [
+    "bigcell_transformer_2d",
+    "BigCellTransformer2DModel"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKLFlux2"
+  ],
+  "scheduler": [
+    "diffusers",
+    "FlowMatchEulerDiscreteScheduler"
+  ]
+}

bigcell-2048/pipeline.py

@@ -0,0 +1,151 @@
+# BigCell pipeline. https://github.com/cvlab-stonybrook/BigCell
+#
+# Generates histopathology images from a grid of UNI-2 embeddings (one 1536-d token per 256x256 px at 20x).
+# A grid that matches the model's native tile (e.g. 8x8 for BigCell-2048) is denoised directly; larger grids
+# are generated with the sliding-window procedure of the paper (BigCell-large), optionally with adaptive
+# window scheduling (``top_r < 1``). For multi-GPU generation of 16k-32k images use the GitHub scripts.
+
+import math
+from typing import List, Optional, Union
+
+import numpy as np
+import torch
+from diffusers import DiffusionPipeline
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.pipelines.pipeline_utils import ImagePipelineOutput
+from diffusers.utils.torch_utils import randn_tensor
+from PIL import Image
+
+LATENTS_PER_TOKEN = 32  # one UNI-2 token covers 256 px = 32 Flux.2 latents
+
+
+def _gaussian_kernel(size, sigma, device):
+    x = torch.linspace(-1, 1, size)
+    x = torch.stack((x.tile(size, 1), x.tile(size, 1).T), dim=0)
+    k = torch.exp(-(torch.linalg.norm(x, dim=0) ** 2) / (2 * sigma**2))
+    return (k / k.max()).to(device)
+
+
+class BigCellPipeline(DiffusionPipeline):
+    model_cpu_offload_seq = "transformer->vae"
+
+    def __init__(self, transformer, vae, scheduler):
+        super().__init__()
+        self.register_modules(transformer=transformer, vae=vae, scheduler=scheduler)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=8)
+
+    @property
+    def tile_tokens(self):
+        """Side of the UNI-2 grid covered by one native tile (1, 4 or 8)."""
+        return int(round(math.sqrt(self.transformer.config.num_caption_tokens)))
+
+    def get_unconditional_embedding(self, batch_size=1):
+        return self.transformer.unconditional_embedding.tile(batch_size, 1, 1)
+
+    def _velocity(self, x, t, y, guidance_scale):
+        """Returns (cfg velocity, importance score |v_cond - v_uncond|)."""
+        null = self.get_unconditional_embedding(x.shape[0]).to(y.dtype)
+        v = self.transformer(torch.cat([x, x]), t.expand(2 * x.shape[0]), torch.cat([y, null])).float()
+        v_cond, v_uncond = v.chunk(2)
+        return v_uncond + guidance_scale * (v_cond - v_uncond), (v_cond - v_uncond).abs().mean().item()
+
+    def _sigmas(self, num_inference_steps, device):
+        self.scheduler.set_timesteps(sigmas=np.linspace(1.0, 0.0, num_inference_steps + 1)[:-1], device=device)
+        return self.scheduler.sigmas.to(device)  # shifted, ends with 0
+
+    @torch.no_grad()
+    def _decode(self, latents, tile=256, sigma=0.5):
+        """Tiled decode for large canvases (Gaussian-blended, accumulated on CPU)."""
+        _, _, h, w = latents.shape
+        dtype = self.vae.dtype
+        if h <= tile and w <= tile:
+            return self.vae.decode(latents.to(dtype)).sample.float().cpu()
+        out = torch.zeros(latents.shape[0], 3, h * 8, w * 8)
+        cnt = torch.zeros(1, 1, h * 8, w * 8)
+        kernel = _gaussian_kernel(tile * 8, sigma, "cpu")[None, None]
+        stride = tile // 2
+        for i in range(0, h - tile + 1, stride):
+            for j in range(0, w - tile + 1, stride):
+                dec = self.vae.decode(latents[:, :, i:i + tile, j:j + tile].to(dtype)).sample.float().cpu()
+                out[:, :, i * 8:(i + tile) * 8, j * 8:(j + tile) * 8] += dec * kernel
+                cnt[:, :, i * 8:(i + tile) * 8, j * 8:(j + tile) * 8] += kernel
+        return out / cnt.clamp(min=1e-6)
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        uni_embeds: torch.Tensor,
+        num_inference_steps: int = 20,
+        guidance_scale: float = 2.0,
+        top_r: float = 1.0,
+        stride: Optional[int] = None,
+        window_sigma: float = 0.5,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+        progress: bool = True,
+    ):
+        """
+        Args:
+            uni_embeds: UNI-2 grid of shape (B, Hg, Wg, 1536), or tokens of one native tile (B, T, 1536).
+                The output image is (256 * Hg) x (256 * Wg) pixels.
+            num_inference_steps: Euler steps (paper: 20).
+            guidance_scale: classifier-free guidance scale (paper: 2.0).
+            top_r: fraction of sliding windows recomputed at every step after the first (paper Sec. 3.3);
+                the rest reuse their cached velocity. 1.0 = full computation.
+            stride: window stride in latents (default: half a window, e.g. 128 = 1024 px for BigCell-2048).
+            window_sigma: width of the Gaussian window weighting.
+            output_type: "pil", "np" (uint8 HxWx3) or "latent".
+        """
+        device = self._execution_device
+        dtype = self.transformer.dtype
+        tok = self.tile_tokens
+        if uni_embeds.ndim == 3:  # (B, T, C) tokens of one tile -> grid
+            uni_embeds = uni_embeds.reshape(uni_embeds.shape[0], tok, tok, -1)
+        B, Hg, Wg, _ = uni_embeds.shape
+        if Hg < tok or Wg < tok:
+            raise ValueError(f"grid {Hg}x{Wg} is smaller than one tile ({tok}x{tok})")
+        uni_embeds = uni_embeds.to(device, dtype)
+
+        win = self.transformer.config.sample_size
+        H, W = Hg * LATENTS_PER_TOKEN, Wg * LATENTS_PER_TOKEN
+        stride = stride or win // 2
+        windows = [(h, w) for h in range(0, H - win + 1, stride) for w in range(0, W - win + 1, stride)]
+        kernel = _gaussian_kernel(win, window_sigma, device)
+
+        if latents is None:
+            latents = randn_tensor((B, self.transformer.config.in_channels, H, W), generator=generator,
+                                   device=device, dtype=torch.float32)
+        x = latents.to(device, torch.float32)
+        sigmas = self._sigmas(num_inference_steps, device)
+
+        n_select = len(windows) if top_r >= 1 else max(1, math.ceil(top_r * len(windows)))
+        importance = [dict() for _ in range(B)]
+        v_prev = torch.zeros_like(x)
+        with self.progress_bar(total=num_inference_steps) as pbar:
+            for step, (s_cur, s_next) in enumerate(zip(sigmas[:-1], sigmas[1:])):
+                t = s_cur.reshape(1).to(dtype)
+                v_new = v_prev.clone()
+                for b in range(B):
+                    active = windows if step == 0 else sorted(windows, key=lambda w: -importance[b].get(w, 0.0))[:n_select]
+                    v_sum = torch.zeros_like(x[b:b + 1])
+                    w_sum = torch.zeros_like(x[b:b + 1])
+                    for h, w in active:
+                        y = uni_embeds[b:b + 1, h // 32:h // 32 + tok, w // 32:w // 32 + tok].reshape(1, tok * tok, -1)
+                        v, score = self._velocity(x[b:b + 1, :, h:h + win, w:w + win].to(dtype), t, y, guidance_scale)
+                        importance[b][(h, w)] = score
+                        v_sum[:, :, h:h + win, w:w + win] += v * kernel
+                        w_sum[:, :, h:h + win, w:w + win] += kernel
+                    v_new[b:b + 1] = torch.where(w_sum > 0, v_sum / w_sum.clamp(min=1e-6), v_prev[b:b + 1])
+                v_prev = v_new
+                x = x + (s_next - s_cur) * v_new
+                if progress:
+                    pbar.update()
+
+        if output_type == "latent":
+            return ImagePipelineOutput(images=x) if return_dict else (x,)
+        img = ((self._decode(x) + 1) / 2).clamp(0, 1)
+        img = (img.permute(0, 2, 3, 1) * 255).round().to(torch.uint8).numpy()
+        images = [Image.fromarray(i) for i in img] if output_type == "pil" else img
+        return ImagePipelineOutput(images=images) if return_dict else (images,)

bigcell-2048/scheduler/scheduler_config.json

@@ -0,0 +1,18 @@
+{
+  "_class_name": "FlowMatchEulerDiscreteScheduler",
+  "_diffusers_version": "0.37.1",
+  "base_image_seq_len": 256,
+  "base_shift": 0.5,
+  "invert_sigmas": false,
+  "max_image_seq_len": 4096,
+  "max_shift": 1.15,
+  "num_train_timesteps": 1000,
+  "shift": 3.0,
+  "shift_terminal": null,
+  "stochastic_sampling": false,
+  "time_shift_type": "exponential",
+  "use_beta_sigmas": false,
+  "use_dynamic_shifting": false,
+  "use_exponential_sigmas": false,
+  "use_karras_sigmas": false
+}

bigcell-2048/transformer/bigcell_transformer_2d.py

@@ -0,0 +1,202 @@
+# BigCell tile generator: a SiT flow-matching transformer conditioned on a grid of UNI-2 embeddings
+# through cross-attention. https://github.com/cvlab-stonybrook/BigCell
+#
+# Adapted from SiT (https://github.com/willisma/SiT, MIT) and REPA-E (https://github.com/End2End-Diffusion/REPA-E, MIT).
+
+import math
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+class PatchEmbed(nn.Module):
+    def __init__(self, patch_size, in_channels, hidden_size):
+        super().__init__()
+        self.proj = nn.Conv2d(in_channels, hidden_size, kernel_size=patch_size, stride=patch_size, bias=True)
+
+    def forward(self, x):
+        return self.proj(x).flatten(2).transpose(1, 2)
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features, out_features=None, approximate="none"):
+        super().__init__()
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = nn.GELU(approximate=approximate)
+        self.fc2 = nn.Linear(hidden_features, out_features or in_features)
+
+    def forward(self, x):
+        return self.fc2(self.act(self.fc1(x)))
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads):
+        super().__init__()
+        self.num_heads = num_heads
+        self.qkv = nn.Linear(dim, dim * 3, bias=True)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        q, k, v = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4).unbind(0)
+        x = F.scaled_dot_product_attention(q, k, v)
+        return self.proj(x.transpose(1, 2).reshape(B, N, C))
+
+
+class CrossAttention(nn.Module):
+    def __init__(self, dim, num_heads):
+        super().__init__()
+        self.num_heads = num_heads
+        self.q_linear = nn.Linear(dim, dim)
+        self.kv_linear = nn.Linear(dim, dim * 2)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x, cond):
+        B, N, C = x.shape
+        h = self.num_heads
+        q = self.q_linear(x).reshape(B, N, h, C // h).transpose(1, 2)
+        k, v = self.kv_linear(cond).reshape(B, cond.shape[1], 2, h, C // h).permute(2, 0, 3, 1, 4).unbind(0)
+        x = F.scaled_dot_product_attention(q, k, v)
+        return self.proj(x.transpose(1, 2).reshape(B, N, C))
+
+
+class TimestepEmbedder(nn.Module):
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    def forward(self, t):
+        half = self.frequency_embedding_size // 2
+        freqs = torch.exp(-math.log(10000) * torch.arange(half, dtype=torch.float32, device=t.device) / half)
+        args = t[:, None].float() * freqs[None]
+        return self.mlp(torch.cat([torch.cos(args), torch.sin(args)], dim=-1).to(t.dtype))
+
+
+class CaptionEmbedder(nn.Module):
+    """Projects UNI-2 tokens; ``y_embedding`` is the learned unconditional embedding used for CFG."""
+
+    def __init__(self, in_channels, hidden_size, token_num):
+        super().__init__()
+        self.y_proj = Mlp(in_channels, hidden_size, hidden_size, approximate="tanh")
+        self.register_buffer("y_embedding", torch.randn(token_num, in_channels) / in_channels**0.5)
+
+    def forward(self, caption):
+        return self.y_proj(caption)
+
+
+class BigCellBlock(nn.Module):
+    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.attn = Attention(hidden_size, num_heads)
+        self.cross_attn = CrossAttention(hidden_size, num_heads)
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.mlp = Mlp(hidden_size, int(hidden_size * mlp_ratio), approximate="tanh")
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True))
+
+    def forward(self, x, cond, t):
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(t).chunk(6, dim=-1)
+        x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa))
+        x = x + self.cross_attn(x, cond)
+        x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
+        return x
+
+
+class FinalLayer(nn.Module):
+    def __init__(self, hidden_size, patch_size, out_channels):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
+
+    def forward(self, x, c):
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)
+        return self.linear(modulate(self.norm_final(x), shift, scale))
+
+
+class BigCellTransformer2DModel(ModelMixin, ConfigMixin):
+    """SiT-XL/2 velocity predictor operating on Flux.2 VAE latents (32 channels, 8x downsampling).
+
+    Args:
+        sample_size: latent side length (tile resolution // 8).
+        num_caption_tokens: number of UNI-2 tokens per tile, (tile resolution // 256) ** 2.
+        alpha_shift: time shift used during training; the scheduler uses the same value.
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        sample_size: int = 256,
+        patch_size: int = 2,
+        in_channels: int = 32,
+        hidden_size: int = 1152,
+        num_layers: int = 28,
+        num_attention_heads: int = 16,
+        mlp_ratio: float = 4.0,
+        caption_channels: int = 1536,
+        num_caption_tokens: int = 64,
+        alpha_shift: float = 1.0,
+    ):
+        super().__init__()
+        self.x_embedder = PatchEmbed(patch_size, in_channels, hidden_size)
+        self.t_embedder = TimestepEmbedder(hidden_size)
+        self.y_embedder = CaptionEmbedder(caption_channels, hidden_size, num_caption_tokens)
+        num_patches = (sample_size // patch_size) ** 2
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, hidden_size), requires_grad=False)
+        grid = int(num_caption_tokens**0.5)
+        if num_caption_tokens > 1 and grid**2 == num_caption_tokens:
+            self.y_pos_embed = nn.Parameter(torch.zeros(1, num_caption_tokens, hidden_size), requires_grad=False)
+        else:
+            self.y_pos_embed = None
+        self.blocks = nn.ModuleList(
+            [BigCellBlock(hidden_size, num_attention_heads, mlp_ratio) for _ in range(num_layers)]
+        )
+        self.final_layer = FinalLayer(hidden_size, patch_size, in_channels)
+        pos = _sincos_2d(hidden_size, np.stack(np.meshgrid(np.arange(sample_size // patch_size),
+                                                          np.arange(sample_size // patch_size)), 0).reshape(2, -1))
+        self.pos_embed.data.copy_(torch.from_numpy(pos).float()[None])
+
+    @property
+    def unconditional_embedding(self):
+        """(1, T, 1536) learned null condition for classifier-free guidance."""
+        return self.y_embedder.y_embedding[None]
+
+    def unpatchify(self, x):
+        c, p = self.config.in_channels, self.config.patch_size
+        h = w = int(x.shape[1] ** 0.5)
+        x = torch.einsum("nhwpqc->nchpwq", x.reshape(x.shape[0], h, w, p, p, c))
+        return x.reshape(x.shape[0], c, h * p, w * p)
+
+    def forward(self, hidden_states, timestep, encoder_hidden_states):
+        """hidden_states: (B, 32, H, W) noisy latent; timestep: (B,) in [0, 1] (1 = pure noise);
+        encoder_hidden_states: (B, T, 1536) UNI-2 tokens in raster order. Returns the velocity (noise - data)."""
+        x = self.x_embedder(hidden_states) + self.pos_embed
+        t = self.t_embedder(timestep)
+        y = self.y_embedder(encoder_hidden_states)
+        if self.y_pos_embed is not None:
+            y = y + self.y_pos_embed
+        for block in self.blocks:
+            x = block(x, y, t)
+        return self.unpatchify(self.final_layer(x, t))
+
+
+def _sincos_2d(embed_dim, grid):
+    def one_d(dim, pos):
+        omega = 1.0 / 10000 ** (np.arange(dim // 2, dtype=np.float64) / (dim / 2.0))
+        out = np.einsum("m,d->md", pos.reshape(-1), omega)
+        return np.concatenate([np.sin(out), np.cos(out)], axis=1)
+
+    return np.concatenate([one_d(embed_dim // 2, grid[0]), one_d(embed_dim // 2, grid[1])], axis=1)

bigcell-2048/transformer/config.json

@@ -0,0 +1,14 @@
+{
+  "_class_name": "BigCellTransformer2DModel",
+  "_diffusers_version": "0.37.1",
+  "alpha_shift": 3.0,
+  "caption_channels": 1536,
+  "hidden_size": 1152,
+  "in_channels": 32,
+  "mlp_ratio": 4.0,
+  "num_attention_heads": 16,
+  "num_caption_tokens": 64,
+  "num_layers": 28,
+  "patch_size": 2,
+  "sample_size": 256
+}

bigcell-2048/transformer/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:17a3b7165206aa7d65140b685ae854a7c96098f2ebe98341371721d2855d030b
+size 3379372944

bigcell-2048/vae/config.json

@@ -0,0 +1,40 @@
+{
+  "_class_name": "AutoencoderKLFlux2",
+  "_diffusers_version": "0.36.0.dev0",
+  "act_fn": "silu",
+  "batch_norm_eps": 0.0001,
+  "batch_norm_momentum": 0.1,
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "force_upcast": true,
+  "in_channels": 3,
+  "latent_channels": 32,
+  "layers_per_block": 2,
+  "mid_block_add_attention": true,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "patch_size": [
+    2,
+    2
+  ],
+  "sample_size": 1024,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ],
+  "use_post_quant_conv": true,
+  "use_quant_conv": true,
+  "_name_or_path": "black-forest-labs/FLUX.2-dev"
+}

bigcell-256/model_index.json

@@ -0,0 +1,19 @@
+{
+  "_class_name": [
+    "pipeline",
+    "BigCellPipeline"
+  ],
+  "_diffusers_version": "0.37.1",
+  "transformer": [
+    "bigcell_transformer_2d",
+    "BigCellTransformer2DModel"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKLFlux2"
+  ],
+  "scheduler": [
+    "diffusers",
+    "FlowMatchEulerDiscreteScheduler"
+  ]
+}

bigcell-256/pipeline.py

@@ -0,0 +1,151 @@
+# BigCell pipeline. https://github.com/cvlab-stonybrook/BigCell
+#
+# Generates histopathology images from a grid of UNI-2 embeddings (one 1536-d token per 256x256 px at 20x).
+# A grid that matches the model's native tile (e.g. 8x8 for BigCell-2048) is denoised directly; larger grids
+# are generated with the sliding-window procedure of the paper (BigCell-large), optionally with adaptive
+# window scheduling (``top_r < 1``). For multi-GPU generation of 16k-32k images use the GitHub scripts.
+
+import math
+from typing import List, Optional, Union
+
+import numpy as np
+import torch
+from diffusers import DiffusionPipeline
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.pipelines.pipeline_utils import ImagePipelineOutput
+from diffusers.utils.torch_utils import randn_tensor
+from PIL import Image
+
+LATENTS_PER_TOKEN = 32  # one UNI-2 token covers 256 px = 32 Flux.2 latents
+
+
+def _gaussian_kernel(size, sigma, device):
+    x = torch.linspace(-1, 1, size)
+    x = torch.stack((x.tile(size, 1), x.tile(size, 1).T), dim=0)
+    k = torch.exp(-(torch.linalg.norm(x, dim=0) ** 2) / (2 * sigma**2))
+    return (k / k.max()).to(device)
+
+
+class BigCellPipeline(DiffusionPipeline):
+    model_cpu_offload_seq = "transformer->vae"
+
+    def __init__(self, transformer, vae, scheduler):
+        super().__init__()
+        self.register_modules(transformer=transformer, vae=vae, scheduler=scheduler)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=8)
+
+    @property
+    def tile_tokens(self):
+        """Side of the UNI-2 grid covered by one native tile (1, 4 or 8)."""
+        return int(round(math.sqrt(self.transformer.config.num_caption_tokens)))
+
+    def get_unconditional_embedding(self, batch_size=1):
+        return self.transformer.unconditional_embedding.tile(batch_size, 1, 1)
+
+    def _velocity(self, x, t, y, guidance_scale):
+        """Returns (cfg velocity, importance score |v_cond - v_uncond|)."""
+        null = self.get_unconditional_embedding(x.shape[0]).to(y.dtype)
+        v = self.transformer(torch.cat([x, x]), t.expand(2 * x.shape[0]), torch.cat([y, null])).float()
+        v_cond, v_uncond = v.chunk(2)
+        return v_uncond + guidance_scale * (v_cond - v_uncond), (v_cond - v_uncond).abs().mean().item()
+
+    def _sigmas(self, num_inference_steps, device):
+        self.scheduler.set_timesteps(sigmas=np.linspace(1.0, 0.0, num_inference_steps + 1)[:-1], device=device)
+        return self.scheduler.sigmas.to(device)  # shifted, ends with 0
+
+    @torch.no_grad()
+    def _decode(self, latents, tile=256, sigma=0.5):
+        """Tiled decode for large canvases (Gaussian-blended, accumulated on CPU)."""
+        _, _, h, w = latents.shape
+        dtype = self.vae.dtype
+        if h <= tile and w <= tile:
+            return self.vae.decode(latents.to(dtype)).sample.float().cpu()
+        out = torch.zeros(latents.shape[0], 3, h * 8, w * 8)
+        cnt = torch.zeros(1, 1, h * 8, w * 8)
+        kernel = _gaussian_kernel(tile * 8, sigma, "cpu")[None, None]
+        stride = tile // 2
+        for i in range(0, h - tile + 1, stride):
+            for j in range(0, w - tile + 1, stride):
+                dec = self.vae.decode(latents[:, :, i:i + tile, j:j + tile].to(dtype)).sample.float().cpu()
+                out[:, :, i * 8:(i + tile) * 8, j * 8:(j + tile) * 8] += dec * kernel
+                cnt[:, :, i * 8:(i + tile) * 8, j * 8:(j + tile) * 8] += kernel
+        return out / cnt.clamp(min=1e-6)
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        uni_embeds: torch.Tensor,
+        num_inference_steps: int = 20,
+        guidance_scale: float = 2.0,
+        top_r: float = 1.0,
+        stride: Optional[int] = None,
+        window_sigma: float = 0.5,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+        progress: bool = True,
+    ):
+        """
+        Args:
+            uni_embeds: UNI-2 grid of shape (B, Hg, Wg, 1536), or tokens of one native tile (B, T, 1536).
+                The output image is (256 * Hg) x (256 * Wg) pixels.
+            num_inference_steps: Euler steps (paper: 20).
+            guidance_scale: classifier-free guidance scale (paper: 2.0).
+            top_r: fraction of sliding windows recomputed at every step after the first (paper Sec. 3.3);
+                the rest reuse their cached velocity. 1.0 = full computation.
+            stride: window stride in latents (default: half a window, e.g. 128 = 1024 px for BigCell-2048).
+            window_sigma: width of the Gaussian window weighting.
+            output_type: "pil", "np" (uint8 HxWx3) or "latent".
+        """
+        device = self._execution_device
+        dtype = self.transformer.dtype
+        tok = self.tile_tokens
+        if uni_embeds.ndim == 3:  # (B, T, C) tokens of one tile -> grid
+            uni_embeds = uni_embeds.reshape(uni_embeds.shape[0], tok, tok, -1)
+        B, Hg, Wg, _ = uni_embeds.shape
+        if Hg < tok or Wg < tok:
+            raise ValueError(f"grid {Hg}x{Wg} is smaller than one tile ({tok}x{tok})")
+        uni_embeds = uni_embeds.to(device, dtype)
+
+        win = self.transformer.config.sample_size
+        H, W = Hg * LATENTS_PER_TOKEN, Wg * LATENTS_PER_TOKEN
+        stride = stride or win // 2
+        windows = [(h, w) for h in range(0, H - win + 1, stride) for w in range(0, W - win + 1, stride)]
+        kernel = _gaussian_kernel(win, window_sigma, device)
+
+        if latents is None:
+            latents = randn_tensor((B, self.transformer.config.in_channels, H, W), generator=generator,
+                                   device=device, dtype=torch.float32)
+        x = latents.to(device, torch.float32)
+        sigmas = self._sigmas(num_inference_steps, device)
+
+        n_select = len(windows) if top_r >= 1 else max(1, math.ceil(top_r * len(windows)))
+        importance = [dict() for _ in range(B)]
+        v_prev = torch.zeros_like(x)
+        with self.progress_bar(total=num_inference_steps) as pbar:
+            for step, (s_cur, s_next) in enumerate(zip(sigmas[:-1], sigmas[1:])):
+                t = s_cur.reshape(1).to(dtype)
+                v_new = v_prev.clone()
+                for b in range(B):
+                    active = windows if step == 0 else sorted(windows, key=lambda w: -importance[b].get(w, 0.0))[:n_select]
+                    v_sum = torch.zeros_like(x[b:b + 1])
+                    w_sum = torch.zeros_like(x[b:b + 1])
+                    for h, w in active:
+                        y = uni_embeds[b:b + 1, h // 32:h // 32 + tok, w // 32:w // 32 + tok].reshape(1, tok * tok, -1)
+                        v, score = self._velocity(x[b:b + 1, :, h:h + win, w:w + win].to(dtype), t, y, guidance_scale)
+                        importance[b][(h, w)] = score
+                        v_sum[:, :, h:h + win, w:w + win] += v * kernel
+                        w_sum[:, :, h:h + win, w:w + win] += kernel
+                    v_new[b:b + 1] = torch.where(w_sum > 0, v_sum / w_sum.clamp(min=1e-6), v_prev[b:b + 1])
+                v_prev = v_new
+                x = x + (s_next - s_cur) * v_new
+                if progress:
+                    pbar.update()
+
+        if output_type == "latent":
+            return ImagePipelineOutput(images=x) if return_dict else (x,)
+        img = ((self._decode(x) + 1) / 2).clamp(0, 1)
+        img = (img.permute(0, 2, 3, 1) * 255).round().to(torch.uint8).numpy()
+        images = [Image.fromarray(i) for i in img] if output_type == "pil" else img
+        return ImagePipelineOutput(images=images) if return_dict else (images,)

bigcell-256/scheduler/scheduler_config.json

@@ -0,0 +1,18 @@
+{
+  "_class_name": "FlowMatchEulerDiscreteScheduler",
+  "_diffusers_version": "0.37.1",
+  "base_image_seq_len": 256,
+  "base_shift": 0.5,
+  "invert_sigmas": false,
+  "max_image_seq_len": 4096,
+  "max_shift": 1.15,
+  "num_train_timesteps": 1000,
+  "shift": 1.0,
+  "shift_terminal": null,
+  "stochastic_sampling": false,
+  "time_shift_type": "exponential",
+  "use_beta_sigmas": false,
+  "use_dynamic_shifting": false,
+  "use_exponential_sigmas": false,
+  "use_karras_sigmas": false
+}

bigcell-256/transformer/bigcell_transformer_2d.py

@@ -0,0 +1,202 @@
+# BigCell tile generator: a SiT flow-matching transformer conditioned on a grid of UNI-2 embeddings
+# through cross-attention. https://github.com/cvlab-stonybrook/BigCell
+#
+# Adapted from SiT (https://github.com/willisma/SiT, MIT) and REPA-E (https://github.com/End2End-Diffusion/REPA-E, MIT).
+
+import math
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+class PatchEmbed(nn.Module):
+    def __init__(self, patch_size, in_channels, hidden_size):
+        super().__init__()
+        self.proj = nn.Conv2d(in_channels, hidden_size, kernel_size=patch_size, stride=patch_size, bias=True)
+
+    def forward(self, x):
+        return self.proj(x).flatten(2).transpose(1, 2)
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features, out_features=None, approximate="none"):
+        super().__init__()
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = nn.GELU(approximate=approximate)
+        self.fc2 = nn.Linear(hidden_features, out_features or in_features)
+
+    def forward(self, x):
+        return self.fc2(self.act(self.fc1(x)))
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads):
+        super().__init__()
+        self.num_heads = num_heads
+        self.qkv = nn.Linear(dim, dim * 3, bias=True)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        q, k, v = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4).unbind(0)
+        x = F.scaled_dot_product_attention(q, k, v)
+        return self.proj(x.transpose(1, 2).reshape(B, N, C))
+
+
+class CrossAttention(nn.Module):
+    def __init__(self, dim, num_heads):
+        super().__init__()
+        self.num_heads = num_heads
+        self.q_linear = nn.Linear(dim, dim)
+        self.kv_linear = nn.Linear(dim, dim * 2)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x, cond):
+        B, N, C = x.shape
+        h = self.num_heads
+        q = self.q_linear(x).reshape(B, N, h, C // h).transpose(1, 2)
+        k, v = self.kv_linear(cond).reshape(B, cond.shape[1], 2, h, C // h).permute(2, 0, 3, 1, 4).unbind(0)
+        x = F.scaled_dot_product_attention(q, k, v)
+        return self.proj(x.transpose(1, 2).reshape(B, N, C))
+
+
+class TimestepEmbedder(nn.Module):
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    def forward(self, t):
+        half = self.frequency_embedding_size // 2
+        freqs = torch.exp(-math.log(10000) * torch.arange(half, dtype=torch.float32, device=t.device) / half)
+        args = t[:, None].float() * freqs[None]
+        return self.mlp(torch.cat([torch.cos(args), torch.sin(args)], dim=-1).to(t.dtype))
+
+
+class CaptionEmbedder(nn.Module):
+    """Projects UNI-2 tokens; ``y_embedding`` is the learned unconditional embedding used for CFG."""
+
+    def __init__(self, in_channels, hidden_size, token_num):
+        super().__init__()
+        self.y_proj = Mlp(in_channels, hidden_size, hidden_size, approximate="tanh")
+        self.register_buffer("y_embedding", torch.randn(token_num, in_channels) / in_channels**0.5)
+
+    def forward(self, caption):
+        return self.y_proj(caption)
+
+
+class BigCellBlock(nn.Module):
+    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.attn = Attention(hidden_size, num_heads)
+        self.cross_attn = CrossAttention(hidden_size, num_heads)
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.mlp = Mlp(hidden_size, int(hidden_size * mlp_ratio), approximate="tanh")
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True))
+
+    def forward(self, x, cond, t):
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(t).chunk(6, dim=-1)
+        x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa))
+        x = x + self.cross_attn(x, cond)
+        x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
+        return x
+
+
+class FinalLayer(nn.Module):
+    def __init__(self, hidden_size, patch_size, out_channels):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
+
+    def forward(self, x, c):
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)
+        return self.linear(modulate(self.norm_final(x), shift, scale))
+
+
+class BigCellTransformer2DModel(ModelMixin, ConfigMixin):
+    """SiT-XL/2 velocity predictor operating on Flux.2 VAE latents (32 channels, 8x downsampling).
+
+    Args:
+        sample_size: latent side length (tile resolution // 8).
+        num_caption_tokens: number of UNI-2 tokens per tile, (tile resolution // 256) ** 2.
+        alpha_shift: time shift used during training; the scheduler uses the same value.
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        sample_size: int = 256,
+        patch_size: int = 2,
+        in_channels: int = 32,
+        hidden_size: int = 1152,
+        num_layers: int = 28,
+        num_attention_heads: int = 16,
+        mlp_ratio: float = 4.0,
+        caption_channels: int = 1536,
+        num_caption_tokens: int = 64,
+        alpha_shift: float = 1.0,
+    ):
+        super().__init__()
+        self.x_embedder = PatchEmbed(patch_size, in_channels, hidden_size)
+        self.t_embedder = TimestepEmbedder(hidden_size)
+        self.y_embedder = CaptionEmbedder(caption_channels, hidden_size, num_caption_tokens)
+        num_patches = (sample_size // patch_size) ** 2
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, hidden_size), requires_grad=False)
+        grid = int(num_caption_tokens**0.5)
+        if num_caption_tokens > 1 and grid**2 == num_caption_tokens:
+            self.y_pos_embed = nn.Parameter(torch.zeros(1, num_caption_tokens, hidden_size), requires_grad=False)
+        else:
+            self.y_pos_embed = None
+        self.blocks = nn.ModuleList(
+            [BigCellBlock(hidden_size, num_attention_heads, mlp_ratio) for _ in range(num_layers)]
+        )
+        self.final_layer = FinalLayer(hidden_size, patch_size, in_channels)
+        pos = _sincos_2d(hidden_size, np.stack(np.meshgrid(np.arange(sample_size // patch_size),
+                                                          np.arange(sample_size // patch_size)), 0).reshape(2, -1))
+        self.pos_embed.data.copy_(torch.from_numpy(pos).float()[None])
+
+    @property
+    def unconditional_embedding(self):
+        """(1, T, 1536) learned null condition for classifier-free guidance."""
+        return self.y_embedder.y_embedding[None]
+
+    def unpatchify(self, x):
+        c, p = self.config.in_channels, self.config.patch_size
+        h = w = int(x.shape[1] ** 0.5)
+        x = torch.einsum("nhwpqc->nchpwq", x.reshape(x.shape[0], h, w, p, p, c))
+        return x.reshape(x.shape[0], c, h * p, w * p)
+
+    def forward(self, hidden_states, timestep, encoder_hidden_states):
+        """hidden_states: (B, 32, H, W) noisy latent; timestep: (B,) in [0, 1] (1 = pure noise);
+        encoder_hidden_states: (B, T, 1536) UNI-2 tokens in raster order. Returns the velocity (noise - data)."""
+        x = self.x_embedder(hidden_states) + self.pos_embed
+        t = self.t_embedder(timestep)
+        y = self.y_embedder(encoder_hidden_states)
+        if self.y_pos_embed is not None:
+            y = y + self.y_pos_embed
+        for block in self.blocks:
+            x = block(x, y, t)
+        return self.unpatchify(self.final_layer(x, t))
+
+
+def _sincos_2d(embed_dim, grid):
+    def one_d(dim, pos):
+        omega = 1.0 / 10000 ** (np.arange(dim // 2, dtype=np.float64) / (dim / 2.0))
+        out = np.einsum("m,d->md", pos.reshape(-1), omega)
+        return np.concatenate([np.sin(out), np.cos(out)], axis=1)
+
+    return np.concatenate([one_d(embed_dim // 2, grid[0]), one_d(embed_dim // 2, grid[1])], axis=1)

bigcell-256/transformer/config.json

@@ -0,0 +1,14 @@
+{
+  "_class_name": "BigCellTransformer2DModel",
+  "_diffusers_version": "0.37.1",
+  "alpha_shift": 1.0,
+  "caption_channels": 1536,
+  "hidden_size": 1152,
+  "in_channels": 32,
+  "mlp_ratio": 4.0,
+  "num_attention_heads": 16,
+  "num_caption_tokens": 1,
+  "num_layers": 28,
+  "patch_size": 2,
+  "sample_size": 32
+}

bigcell-256/transformer/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:091870c776997dd5314f8ccb69fdac97be13e54813eb8af2e56104cc437d0c78
+size 3304373040

bigcell-256/vae/config.json

@@ -0,0 +1,40 @@
+{
+  "_class_name": "AutoencoderKLFlux2",
+  "_diffusers_version": "0.36.0.dev0",
+  "act_fn": "silu",
+  "batch_norm_eps": 0.0001,
+  "batch_norm_momentum": 0.1,
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "force_upcast": true,
+  "in_channels": 3,
+  "latent_channels": 32,
+  "layers_per_block": 2,
+  "mid_block_add_attention": true,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "patch_size": [
+    2,
+    2
+  ],
+  "sample_size": 1024,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ],
+  "use_post_quant_conv": true,
+  "use_quant_conv": true,
+  "_name_or_path": "black-forest-labs/FLUX.2-dev"
+}

cgm-8192/config.json

@@ -0,0 +1,9 @@
+{
+  "model": "CGM-B",
+  "grid_size": 32,
+  "region_size": 8192,
+  "time_shift": 19.595917942265423,
+  "text_encoder": "MahmoodLab/TITAN",
+  "num_inference_steps": 50,
+  "guidance_scale": 2.0
+}

cgm-8192/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a3ae9348c830c1b0f8df002341ebd919e8ae76c03f5706d1971cafdcc7c22a78
+size 1308806832

cgm-8192/normalization_stats.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:64d6310d55153ff074537018252b4809b80cfd3044354793fff554ad945137b9
+size 12583072