stmdit-anon/stmdit-checkpoints

STMDiT — Inference Checkpoints

EMA-only inference weights for every model row reported in the ICML 2026 SD4H workshop submission Transcriptomics-Conditioned Virtual Tissue Synthesis via Diffusion Transformers.

Each subfolder contains a single model:

• model.pt — EMA-only state dict (PyTorch .pt)
• training_config.yaml — the original training YAML
• README.md — per-model card

Models

row_id	Paper label
pixcell-b	PixCell-B
pixcell-flow-b	PixCell-Flow-B
adaln-ddpm-p01	PixCell-GE-B (p=0.1)
adaln-ddpm-p02	PixCell-GE-B-p02
adaln-ddpm-p03	PixCell-GE-B-p03
adaln-ddpm-p05	PixCell-GE-B-p05
adaln-ddpm-p06	PixCell-GE-B-p06
adaln-flow-p01	PixCell-Flow-GE-B (p=0.1)
adaln-flow-p02	PixCell-Flow-GE-B-p02
adaln-flow-p03	PixCell-Flow-GE-B-p03
adaln-flow-p05	PixCell-Flow-GE-B-p05
xattn-direct-p01	XAttn-Direct (p=0.1)
xattn-gsa-p01	XAttn-GSA (p=0.1)
xattn-perceiver-p01	XAttn-Perceiver (p=0.1)
xattn-pma-p01	XAttn-PMA (p=0.1)
xattn-perceiver-p05	XAttn-Perceiver-p05
xattn-perceiver-p06	XAttn-Perceiver-p06
xattn-pma-p05	XAttn-PMA-p05
xattn-pma-p06	XAttn-PMA-p06
ptpl-adaln-p05	PTPL-AdaLN-B (p=0.5)
ptpl-adaln-p06	PTPL-AdaLN-B-p06
ptpl-adaln-p07	PTPL-AdaLN-B-p07
ptpl-xattn-perceiver-p05	PTPL-XAttn-Perceiver-B (p=0.5)
ptpl-xattn-perceiver-p06	PTPL-XAttn-Perceiver-B-p06
ptpl-xattn-perceiver-p07	PTPL-XAttn-Perceiver-B-p07
ptpl-xattn-pma-p05	PTPL-XAttn-PMA-B (p=0.5)
ptpl-xattn-pma-p06	PTPL-XAttn-PMA-B-p06
ptpl-xattn-pma-p07	PTPL-XAttn-PMA-B-p07

Usage

from huggingface_hub import snapshot_download
ckpt_dir = snapshot_download(
    repo_id="stmdit-anon/stmdit-checkpoints",
    allow_patterns="xattn-perceiver-p05/*",   # pick one model by row_id
)
# ckpt_dir / "xattn-perceiver-p05" / "model.pt"

See the public anonymized code repo (linked from the OpenReview submission) for the loader, demo notebook, and end-to-end inference example.