| { |
| "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" |
| } |
|
|
