Open-X-Cell

Single-cell perturbation modeling · Mini checkpoint

Part of the UPM project.

🤗HuggingFace | GitHub

Overview  ·  Quick start  ·  HEK293T example  ·  Python API  ·  Training  ·  Model details

Training and inference code for X-Cell, with model checkpoints trained on public Orion data.

Overview

Open-X-Cell is part of the UPM (Unified Perturbation Modeling) project, supporting the work How Much Can AI Virtual Cells Gain from Unified Perturbational Modeling? This repository provides the X-Cell implementation and checkpoint release; visit the UPM main repository for the overall project.

Open-X-Cell predicts single-cell gene expression after a genetic perturbation, starting from control cells. The Mini model combines control expression, a target gene and five fixed biological priors, then predicts a response through four diffusion steps.

60.0M parameters Trained Mini checkpoint

19,408 genes Checkpoint vocabulary

5 biological priors Protein · network · cellular features

Component	What is provided
Genetic perturbation prediction	One target gene and a pool of control cells → predicted perturbed expression.
Model checkpoint	Original checkpoint.pt, portable Safetensors weights, configuration, vocabulary and all five priors.
Training and inference code	Training from scratch, fine-tuning, exact training resume, prediction CLI and Python API in the GitHub repository.

This is an independent public-data reproduction, not Xaira's official implementation or Pisces checkpoint. The release supports single-target genetic perturbations; it does not provide drug, cytokine or combination prediction interfaces.

Quick start

1. Install

Use Python ≥ 3.10 and PyTorch ≥ 2.3. For GPU inference, install a CUDA-compatible PyTorch build first.

git clone https://github.com/zongtingwei/Open-X-cell.git
cd Open-X-cell
python -m pip install -e .

The checkout includes the HEK293T example used below. Both the code repository and this HF repository currently require authorized access.

2. Download the complete checkpoint

hf auth login
open-x-cell-download --repo zongting/Open-X-Cell   --revision 21b740c097b839eb413749fa6a5a7102d4d2b67d   --output checkpoints/mini-orion

This downloads the nine-file inference bundle, approximately 756 MB, and checks each model file against its size and SHA-256 in the manifest. Allow 3 GB of disk headroom for downloads and example outputs. Choose a new destination directory; the downloader refuses to overwrite an existing one.

Checkpoint contents · complete inference bundle + original PyTorch export

File	Purpose
model.safetensors	Float32 inference weights, approximately 240 MB.
config.json	Model architecture and dimensions.
genes.json	Ordered vocabulary of 19,408 gene symbols.
priors/	Five fixed-prior arrays, approximately 516 MB in total.
manifest.json	Checksums, training settings, provenance and prior coverage.
checkpoint.pt	Original PyTorch training export, approximately 240 MB; available separately.

The weights and all five priors are required for the inference interface. The original PyTorch checkpoint and the Safetensors export contain identical model weights. To download only the original checkpoint:

hf download zongting/Open-X-Cell checkpoint.pt   --revision 21b740c097b839eb413749fa6a5a7102d4d2b67d   --local-dir checkpoints/original

checkpoint.pt contains model weights, configuration and vocabulary. Neither export contains optimizer or scheduler state; exact training resume uses a training run's latest.pt.

3. Run the HEK293T example

open-x-cell-infer   --model checkpoints/mini-orion   --control examples/hek293t/controls.h5ad   --perturbation STAT1   --input-space log1p-cp10k   --device cpu   --n-cells 32   --max-genes 256   --seed 0   --output outputs/STAT1-smoke.h5ad

The output contains 32 predicted cells × 256 genes. This reduced gene context keeps the example small; it changes the model's context and predictions. For the full overlapping gene panel, omit --max-genes and plan for greater memory use. GPU inference supports --device cuda --precision bfloat16 on compatible hardware. Choose a new output filename for each run.

HEK293T example

A compact example from the QC-complete X-Atlas-Orion HEK293T data is included in examples/hek293t. It contains 32 non-targeting controls and 32 separately sampled STAT1-perturbed cells.

File	Cells	Role
controls.h5ad	32	Starting control states used as inference input.
observed.h5ad	32	Experimentally observed STAT1 perturbation cells for inspection.
provenance.json	—	Source rows, seed 0 sampling, expression space and checksums.

Both H5AD files retain all 18,997 measured genes. Expression in .X is already log1p(CP10K), and .var_names contains gene symbols. Control and observed cells are disjoint.

Expected smoke-test output: 32 controls + one STAT1 request → 32 predictions across 256 selected genes. Inference reads only controls.h5ad; the observed cells are provided for inspection.

HEK293T was part of training, so this example is not an independent held-out benchmark. Predicted and observed cells are not paired cell by cell. Data source: X-Atlas-Orion.

Python API

Load the downloaded bundle and run the same example from Python:

from open_x_cell import OpenXCell

model = OpenXCell.from_pretrained("checkpoints/mini-orion", device="cpu")
prediction = model.predict(
    "examples/hek293t/controls.h5ad",
    perturbation="STAT1",
    output="outputs/STAT1-python.h5ad",
    input_space="log1p-cp10k",
    n_cells=32,
    max_genes=256,
    seed=0,
)
print(prediction.shape)  # (32, 256)

from_pretrained accepts a local bundle directory. Download from HF first using open-x-cell-download; each prediction loads the weights and fixed priors from that bundle.

Training and fine-tuning

Both training and inference implementations are included in the GitHub repository. Training requires control cells and observed perturbation cells, with perturbation labels and a compatible gene panel. Use the repository's data preparation instructions before starting a run.

Fine-tune the released Mini checkpoint

The following example expects pre-normalized training expression and a compatible vocabulary:

open-x-cell train \
  --adata /path/to/training.h5ad --output runs/mini-finetune \
  --model-size mini --normalize none --device cuda \
  --pretrained-checkpoint checkpoints/mini-orion/model.safetensors \
  --vocabulary checkpoints/mini-orion/genes.json \
  --esm-priors checkpoints/mini-orion/priors/esm.npz \
  --string-priors checkpoints/mini-orion/priors/string.npz \
  --genept-priors checkpoints/mini-orion/priors/genept.npz \
  --depmap-priors checkpoints/mini-orion/priors/depmap.npz \
  --cell-painting-priors checkpoints/mini-orion/priors/cell_painting.npz \
  --training-context-length 4000 \
  --batch-size 1 --set-size 64 --max-steps 1000 \
  --learning-rate 0.0001 --precision bfloat16 \
  --gradient-accumulation-steps 8 --gradient-checkpointing \
  --checkpoint-every 50

Fine-tuning starts a fresh optimizer and scheduler. Omit --pretrained-checkpoint to train from scratch. These are example settings, not Xaira's private training protocol. Choose context length and batch size for your hardware.

A training run writes checkpoint.pt for model weights and latest.pt for model, optimizer and scheduler state. To resume an interrupted run, keep the original data, priors, settings and training horizon, remove initialization flags, and add --resume /path/to/latest.pt. See training and resume instructions for the complete command requirements.

Model details

Property	This checkpoint
Model	Open-X-Cell Mini · public Orion proxy
Parameters	60,001,618
Vocabulary	19,408 gene symbols
Transformer	12 layers · hidden width 512 · 8 heads
Diffusion steps	4
Training contexts	Public Orion HCT116 and HEK293T
Training budget	4,155 optimizer steps · approximately 6 effective epochs
Training gene contexts	Dynamic, target-first, 4,000 genes
Numerical format	Float32 Safetensors weights and fixed-prior arrays
Input	Control-cell expression and one target gene
Output	Predicted expression over aligned input genes
Hardware	CPU or CUDA

Fixed priors · feature dimensions and vocabulary coverage

Source	Dimensions	Available checkpoint genes
ESM-2	5,120	18,904
STRING	512	18,851
GenePT	3,072	18,579
DepMap	1,150	17,334
Cell Painting	259	7,851

The bundle retains available source rows for the model vocabulary. Missing gene/source pairs use the model's existing missing-prior behavior. Coverage counts describe available features, not experimentally validated perturbations.

Using your own cells

• Supply a control-only .h5ad and a target gene from the checkpoint vocabulary. If .obs['is_control'] is present, all rows must be controls.
• For raw counts, use --input-space counts; add --layer counts if counts are stored in that layer. Normalization uses the complete supplied measured panel before gene alignment.
• For normalized expression, use --input-space log1p-cp10k.
• Gene symbols in .var_names work directly. Ensembl identifiers can be accompanied by supported alias columns such as .var['gene_name'] or .var['gene_symbol'].
• --n-cells sets the number of controls per set. One set uses the first available controls and wrap-pads a short pool; --n-sets greater than one samples controls using the supplied seed.

Validation and scope

The release passed 104 software tests, including export integrity, corrupted-download rejection, input validation, training and inference paths. Installed-package checks also verified training resume against uninterrupted training. The original checkpoint and portable export produced identical predictions in a deterministic CPU smoke test; downloading the bundle from HF reproduced those predictions exactly (32 cells × 256 genes, STAT1).

These checks establish software and export consistency, not biological accuracy. This model does not include private Pisces screens, original private vocabulary/remap artifacts or unpublished QC choices. Mini is the released trained checkpoint; Tiny is intended for smoke tests, and no trained Ultra checkpoint is included. Detailed training settings and limitations are recorded in manifest.json.

Licensing and attribution

The Open-X-Cell model checkpoints and source code are licensed under CC BY-NC 4.0 (Creative Commons Attribution-NonCommercial 4.0 International). Third-party data and fixed-prior assets retain their upstream terms and attribution: X-Atlas-Orion, ESM, STRING, GenePT, DepMap, and JUMP Cell Painting.

Paper screenshot: first-page title and author crop from Wang et al. (2026), under CC BY-NC 4.0.

Citation

If you use this code, please cite our paper.

@misc{upm2026,
  title={How Much Can AI Virtual Cells Gain from Unified Perturbational Modeling?},
  year={2026},
  note={Manuscript in preparation},
  url={https://github.com/youlab2025/UPM}
}

For the original X-Cell method, please cite:

@article{wang2026x,
  title={X-cell: Scaling causal perturbation prediction across diverse cellular contexts via diffusion language models},
  author={Wang, Chloe and Karimzadeh, Mehran and Ravindra, Neal G and Bounds, Lexi R and Alerasool, Nader and Huang, Ann C and Ma, Shihao and Gulbranson, Daniel R and Cui, Haotian and Lee, Yongju and others},
  journal={bioRxiv},
  pages={2026--03},
  year={2026},
  publisher={Cold Spring Harbor Laboratory}
}