SALA

Symmetry-Aware Lattice and Asymmetric-unit generation

Generate molecular crystal candidates with a symmetry-constrained flow model. This repository hosts the full-precision pretrained checkpoint; preprocessing, training and inference code are available on GitHub.

Code and quick start · Inference guide · Project website

Start from SMILES

git clone https://github.com/WENDI-CAI/SALA.git
cd SALA
python -m pip install -e '.[inference]'

sala-predict \
  --smiles 'CC(=O)Oc1ccccc1C(=O)O' \
  --space-group 1 \
  --num-samples 4 \
  --output outputs/aspirin

The first run downloads and verifies sala_ema.ckpt (1.95 GB). Later runs reuse the local cache. Input molecules are processed locally. --device cpu selects CPU execution; compatible CUDA devices are used automatically when available.

Space group 1 (P1) is an explicit assumption in this example. SALA does not infer the space group or asymmetric-unit composition from SMILES. Specify the intended --space-group, optionally an exact --hall-number, and --copies of the complete molecular composition in the asymmetric unit. Include counterions and coformers in dot-separated SMILES where appropriate.

Use a local CIF or checkpoint

sala-predict --input /path/to/target.cif --output outputs/crystal

sala-predict \
  --smiles 'CCO' --space-group 1 \
  --checkpoint /path/to/sala_ema.ckpt \
  --device cpu --output outputs/local

The local checkpoint option supports offline execution. --local-files-only uses an already populated Hugging Face cache without downloading. The advanced sala-generate command supports folders and preprocessed dataset indices.

Model specification

Property	Value
Architecture	Full-cell Crystal DiT
Trainable parameters	486,282,617
Transformer	24 layers · 1,024 hidden width · 16 heads
SwiGLU branch width	2,816
Self-conditioning / pair bias	Enabled
Sampler	Flow-2M · 24 steps
Stored precision	float32, no quantization
Weight selection	Power-function EMA, relative standard deviation 0.05
Source optimizer step	249,000
Source epoch	175
Checkpoint schema	16

Training combines full-cell coordinate velocity, lattice velocity, and contact/clash objectives. Coordinate and lattice endpoint-loss weights are zero. config.json records the architecture and flow settings; model.json records the weight provenance and checksum.

The checkpoint contains the selected EMA state, model buffers and public model configuration. It excludes the optimizer, training loop, extra EMA profiles and private data paths. Use it for inference or initialize a new run with finetune_from; exact training continuation requires a full training checkpoint from your own run.

Outputs and scope

Each run writes heavy-atom CIF candidates, generated_results.pt, input metadata and a run manifest. Candidates are unranked and unrelaxed. The model assumes general-position symmetry orbits. Explicit stereochemistry, isotope labels and radicals are outside the supported SMILES representation and are rejected rather than silently discarded. See the code repository for detailed input constraints.

These weights are not served by a Hugging Face hosted inference endpoint. Run the SALA package locally. Memory use depends on the expanded cell atom count and sampling batch size; the checkpoint size is not a complete memory estimate.

Data and license

Training data, crystallographic records, dataset indices and samples are not distributed because of data licensing restrictions. Supply your own appropriately licensed inputs.

The code repository uses the MIT License. A separate license for these model weights has not been specified.

Integrity

SHA-256 for sala_ema.ckpt:

453db9f2dc8af7047d117b156ff3e8568789fbff349f2f3063577620b290804f

The exported tensors were checked for exact equality with the selected source EMA state. Local CPU execution checks cover preprocessing, training/continuation, and full-model sampling; they do not constitute a benchmark or a guarantee of physical validity.