stmdit-checkpoints / README.md
stmdit-anon's picture
Add umbrella README
91758a1 verified
metadata
license: apache-2.0
tags:
  - histopathology
  - diffusion
  - spatial-transcriptomics
  - icml-2026-sd4h-workshop

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.