ZihengZ/SAE-virchow-v2

SAE + Emb2Img for Virchow2 (PICASSO reproduction)

Two checkpoints that reproduce the models of PICASSO (Kim et al., bioRxiv 2026) on TCGA H&E slides:

Folder	What it is	Size
sae/	TopK sparse autoencoder on the Virchow2 CLS token: 1280 → 40,960 concepts, k = 30	0.4 GB
emb2img/	Latent diffusion model (Stable Diffusion 1.5 U-Net) that generates a 256 px H&E tile from a Virchow2 CLS embedding	3.4 GB

Together they let you raise or lower one concept in an embedding and see how the tissue changes. Code: https://github.com/Pediatric-Brain-Tumor/PICASSO_reproduce.

Which model, which layer, which input

Foundation model	paige-ai/Virchow2 via timm (hf-hub:paige-ai/Virchow2, mlp_layer=SwiGLUPacked, act_layer=SiLU)
Token	CLS token of the final layer: model(x)[:, 0], 1280-d, float32, no normalisation
Tile	256 × 256 px at 0.5 µm/px (20×), resized to 224 × 224 (bicubic), then Virchow2's own mean/std
SAE	d_model = 1280, d_sae = 40960, top_k = 30, encoder subtracts b_dec before the top-k (SAELens convention)
Emb2Img	SD 1.5 U-Net at 256 px (32 × 32 latents), VAE stabilityai/sd-vae-ft-mse, PNDM scheduler, 50 steps, guidance 4.5

Other tokens (patch mean, register tokens), other layers, other magnifications or other models are not supported by these weights.

Install

git clone https://github.com/Pediatric-Brain-Tumor/PICASSO_reproduce && cd PICASSO_reproduce
pip install -e .
huggingface-cli login          # Virchow2 is gated: accept its terms first
huggingface-cli download ZihengZ/SAE-virchow-v2 --local-dir ckpt

1. Embed a tile with Virchow2

import torch, numpy as np
from PIL import Image
from picasso.embed import build_model
from picasso.emb2img.sample_eval import virchow_cls

dev = "cuda"
vit, mean, std, _ = build_model("hf-hub:paige-ai/Virchow2", dev)
tile = Image.open("tile.png").convert("RGB").resize((256, 256), Image.LANCZOS)   # 256 px at 0.5 µm/px
x = torch.from_numpy(np.asarray(tile)).permute(2, 0, 1)[None].float().div(255).to(dev)
cls = virchow_cls(vit, mean, std, x)        # [1, 1280] final-layer CLS token

2. SAE: concepts of an embedding

from picasso.model import load as load_sae

sae = load_sae("ckpt/sae", dev)              # reads sae.safetensors + config.json
with torch.inference_mode():
    f = sae.encode(cls)                      # [1, 40960], exactly 30 non-zeros
    recon = sae.decode(f)[:, -1]             # [1, 1280]
top = f[0].topk(10)
print([(int(i), round(float(v), 2)) for v, i in zip(top.values, top.indices)])

sae.W_dec[j] (unit norm, shape [1280]) is the direction of concept j in Virchow2 space; sae.W_enc[:, j] is its read-out direction. If you only need the weights:

from safetensors.torch import load_file
w = load_file("ckpt/sae/sae.safetensors")    # W_enc [1280, 40960], b_enc [40960], W_dec [40960, 1280], b_dec [1280]

Test-set quality (984k held-out TCGA tiles): explained variance 0.83, cosine 0.93, 1,448 dead concepts (sae/test_metrics.json).

3. Emb2Img: embedding → H&E tile

import pathlib
from diffusers import AutoencoderKL, PNDMScheduler
from picasso.emb2img.model import Emb2Img
from picasso.emb2img.sample_eval import SampleConfig, generate

cfg = SampleConfig(checkpoint=pathlib.Path("ckpt/emb2img"), pairs_dir=pathlib.Path("."), output_dir=pathlib.Path("."))
model = Emb2Img.load(cfg.checkpoint, dev)    # unet/ + cond.safetensors
vae = AutoencoderKL.from_pretrained(cfg.vae, torch_dtype=torch.float16).to(dev).eval()
sched = PNDMScheduler.from_pretrained(cfg.sd_model, subfolder="scheduler")

g = torch.Generator(device=dev).manual_seed(0)
img = generate(model, vae, sched, cls, cfg, g)                       # [1, 3, 256, 256] in [0, 1]
Image.fromarray((img[0].permute(1, 2, 0).cpu().numpy() * 255).round().astype(np.uint8)).save("gen.png")

Closed-loop check (embed the generated tile again and compare with the input embedding): cosine 0.75 on average, against 0.72 for the nearest real tile.

4. Concept intervention

Add s · W_dec[j] to the embedding and generate:

for s in [0, 10, 20, 30]:
    cond = cls + s * sae.W_dec.data[j]
    img = generate(model, vae, sched, cond, cfg, torch.Generator(device=dev).manual_seed(0))

Or in one command:

python scripts/demo.py --image tile.png --sae ckpt/sae --emb2img ckpt/emb2img --latent <concept id> --levels 0,10,20,30

Files

sae/sae.safetensors           W_enc, b_enc, W_dec, b_dec
sae/config.json               d_model, d_sae, top_k, input description
sae/test_metrics.json
emb2img/unet/                 diffusers UNet2DConditionModel (SD 1.5 architecture, 768-d cross-attention)
emb2img/cond.safetensors      LayerNorm + Linear(1280 → 768) projection of the embedding, and the null token for guidance

License

Virchow2 weights are CC-BY-NC-ND 4.0 and Stable Diffusion 1.5 is CreativeML OpenRAIL-M; these checkpoints inherit those terms and are for non-commercial research use only. TCGA data are open access through the GDC.