Translation of Protenix to JAX/Equinox. This is pretty rough, we suggest using it through mosaic.
PyTorch is an optional dependency. The full inference pipeline (featurization, model loading, structure prediction) runs without it.
# Inference only (no PyTorch)
uv sync
# With PyTorch (needed for converting checkpoints from the original Protenix format)
uv sync --extra torchPre-converted Equinox models skip the PyTorch dependency entirely and load in under a second.
| Model | Params | .eqx size |
|---|---|---|
protenix_tiny_default_v0.5.0 |
110M | 438 MB |
protenix_mini_default_v0.5.0 |
134M | 536 MB |
protenix_base_default_v1.0.0 |
368M | 1474 MB |
protenix_base_20250630_v1.0.0 |
368M | 1474 MB |
Models are hosted on HuggingFace and downloaded automatically on first use.
from protenix.backend import load_model
# Downloads from HuggingFace, caches to ~/.protenix/
model = load_model("protenix_base_default_v1.0.0")
# Or load from an explicit path
model = load_model("~/.protenix/protenix_base_default_v1.0.0")Requires the torch extra.
uv sync --extra torch
python translate_models.pyThis downloads any missing checkpoints, converts each model to Equinox, saves .eqx + .skeleton.pkl to ~/.protenix/, and verifies a bit-exact round-trip.