You can deploy this model and start earning money today!

josephjojoe/
chai-mlx

Name: chai-mlx
Availability: InStock

mlx

safetensors

protein-structure-prediction

Views

No views yet

chai-mlx weights

MLX-format safetensors weights for Chai-1, converted from the released TorchScript distribution for use with chai-mlx.

Files

File	Contents	Size
`config.json`	Serialized `ChaiConfig`	2 KB
`model.safetensors.index.json`	Sharded weight map	210 KB
`model-trunk.safetensors`	Pairformer trunk	680 MB
`model-diffusion_module.safetensors`	Diffusion module	512 MB
`model-confidence_head.safetensors`	Confidence head	59 MB
`model-token_embedder.safetensors`	Token input embedder	6.6 MB
`model-feature_embedding.safetensors`	Feature embedding stack	4.8 MB
`model-bond_loss_input_proj.safetensors`	Bond feature projection	2 KB

Total size: ~1.2 GB.

Loading

python

from chai_mlx import ChaiMLX

model = ChaiMLX.from_pretrained("josephjojoe/chai-mlx")
model_fp32 = ChaiMLX.from_pretrained(
    "josephjojoe/chai-mlx",
    compute_dtype="float32",
)

config.json sets:

config_version = "1"
compute_dtype = "reference"

reference is the default runtime precision policy in chai-mlx. Use compute_dtype="float32" to keep the MLX port in fp32 throughout.

Provenance

These weights are a 1:1 tensor rename of Chai Discovery's released Chai-1 TorchScript checkpoints into MLX-compatible safetensors. No retraining, finetuning, or numerical modification was performed.

License

Apache-2.0, inherited from upstream Chai-1. See NOTICE in the chai-mlx repo for attribution.