Upload 5 files

README.md

@@ -0,0 +1,92 @@
+---
+language:
+  - en
+library_name: pytorch
+tags:
+  - chemistry
+  - molecular-crystals
+  - crystal-structure-prediction
+  - flow-matching
+  - materials-science
+---
+
+# 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](https://github.com/WENDI-CAI/SALA).
+
+[Code and quick start](https://github.com/WENDI-CAI/SALA) · [Inference guide](https://github.com/WENDI-CAI/SALA/blob/main/docs/inference.md) · [Project website](https://wendi-cai.github.io/SALA/)
+
+## Start from SMILES
+
+```bash
+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
+
+```bash
+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`:
+
+```text
+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.

SHA256SUMS

@@ -0,0 +1 @@
+453db9f2dc8af7047d117b156ff3e8568789fbff349f2f3063577620b290804f  sala_ema.ckpt

config.json

@@ -0,0 +1,146 @@
+{
+  "model_family": "sala_full_cell_crystal_dit_rectified_flow",
+  "checkpoint_schema_version": 16,
+  "denoiser_config": {
+    "architecture": "full_cell_crystal_dit_flow",
+    "d_model": 1024,
+    "num_heads": 16,
+    "num_layers": 24,
+    "d_ffn": 2816,
+    "dropout": 0.0,
+    "num_elements": 119,
+    "use_gradient_checkpointing": false,
+    "embedding": {
+      "num_fourier_freqs": 4,
+      "num_hybridizations": 8,
+      "formal_charge_range": 16,
+      "num_covalent_layers": 2
+    },
+    "attention": {
+      "attn_dropout": 0.0,
+      "use_flash": true,
+      "backend": "sdpa",
+      "qk_norm": true
+    },
+    "pair": {
+      "enabled": true,
+      "d_pair": 128,
+      "num_rbf": 32,
+      "rbf_cutoff": 12.0,
+      "inter_component_cutoff": 8.0,
+      "max_bond_types": 16,
+      "max_image_candidates": 65536,
+      "feature_chunk_size": 131072,
+      "pair_chunk_size": 2048,
+      "image_chunk_size": 256,
+      "cuda_static_image_radius": 2,
+      "cuda_dynamic_image_radius": true
+    },
+    "orbit": {
+      "require_general_position": true,
+      "special_position_tolerance": 1e-05
+    },
+    "self_conditioning": {
+      "enabled": true
+    },
+    "atom_coord_repr": "canonical_cartesian",
+    "lattice_parameterization": "constrained_q",
+    "component_representation": "base_instance_type_v2"
+  },
+  "flow_config": {
+    "state": {
+      "coordinate": "canonical_cartesian",
+      "lattice": "constrained_q",
+      "model": "normalized_cartesian_whitened_q",
+      "output": "fractional_unwrapped",
+      "coord_scale": 8.0
+    },
+    "path": {
+      "type": "rectified_linear",
+      "prediction": "velocity",
+      "data_time": 0.0,
+      "noise_time": 1.0
+    },
+    "time_sampling": {
+      "type": "logit_normal",
+      "mean": 0.0,
+      "std": 1.0,
+      "clamp": 1e-05
+    },
+    "time_condition": {
+      "type": "scaled_time",
+      "scale": 1000.0
+    },
+    "geometry": {
+      "stabilization_start": 0.5,
+      "latent_clip": 1.5
+    },
+    "self_conditioning": {
+      "enabled": true,
+      "train_probability": 0.5,
+      "prepass_selected_only": true,
+      "detach": true,
+      "sampling_enabled": true,
+      "representation": "normalized_clean_residual"
+    },
+    "loss": {
+      "coord_velocity_weight": 1.0,
+      "coord_velocity_mode": "full_cell_orbit",
+      "lattice_velocity_weight": 1.0,
+      "full_cell_endpoint_weight": 0.0,
+      "coord_endpoint_pred_lattice_mode": "input",
+      "lattice_endpoint_weight": 0.0,
+      "lattice_normalize_by_dof": true,
+      "contact_clash": {
+        "enabled": true,
+        "weight": 0.1,
+        "max_time": 0.4,
+        "vdw_offset_A": 1.0,
+        "contact_cutoff_A": 5.0,
+        "contact_match_weight": 1.0,
+        "huber_delta_A": 0.5,
+        "predicted_lattice_mode": "input",
+        "pair_chunk_size": 65536,
+        "checkpoint_pair_distances": false,
+        "periodic_image_radius": 1,
+        "periodic_fallback_max_candidates": 65536
+      }
+    },
+    "prior": {
+      "lattice": {
+        "type": "volume",
+        "length_log_std": 0.3,
+        "angle_std": 0.2
+      }
+    },
+    "sampler": {
+      "solver": "flow_2m",
+      "num_steps": 24,
+      "time_grid": "uniform",
+      "count_nfe": true
+    },
+    "uncertainty_weighting": {
+      "enabled": false,
+      "channels": 128,
+      "clamp": [
+        -4.0,
+        4.0
+      ]
+    },
+    "diagnostics": {
+      "time_bucket_edges": [
+        0.0,
+        0.1,
+        0.25,
+        0.5,
+        0.75,
+        0.9,
+        1.0
+      ]
+    }
+  },
+  "architectures": [
+    "FullCellCrystalDiT"
+  ],
+  "model_type": "sala"
+}

model.json

@@ -0,0 +1,17 @@
+{
+  "format_version": 1,
+  "weights": "ema",
+  "ema_profile_std": 0.05,
+  "source_global_step": 249000,
+  "source_epoch": 175,
+  "trainable_parameters": 486282617,
+  "training_state_included": false,
+  "release_date": "2026-10-21",
+  "filename": "sala_ema.ckpt",
+  "bytes": 1945255862,
+  "sha256": "453db9f2dc8af7047d117b156ff3e8568789fbff349f2f3063577620b290804f",
+  "tensor_count": 335,
+  "floating_point_dtype": "float32",
+  "model_family": "sala_full_cell_crystal_dit_rectified_flow",
+  "checkpoint_schema_version": 16
+}

sala_ema.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:453db9f2dc8af7047d117b156ff3e8568789fbff349f2f3063577620b290804f
+size 1945255862