OneScience-Sugon/La-Proteina

La-Proteina

Model Overview

La-Proteina is a protein structure generation model based on Partially Latent Flow Matching. It can directly generate all-atom protein structures together with their corresponding amino acid sequences.

Paper: La-Proteina: Atomistic Protein Generation via Partially Latent Flow Matching (arXiv 2025).

• Paper link

Model Description

La-Proteina explicitly models the protein backbone (backbone CA), while sequence and atom-level details are captured through fixed-dimensional latent variables for each residue. This effectively avoids the challenges introduced by explicit side-chain representations.

Use Cases

Use case	Description
Protein structure generation	Generates protein backbones (backbone CA) and local latents based on flow matching.
Motif-constrained generation	Supports motif position and sequence constraints for backbone design around functional motifs.
Diffusion model training	Trains the main La-Proteina model on PDB datasets.
Autoencoder training and inference	Trains the local-latent autoencoder and performs encoding, decoding, and reconstruction evaluation on PDB structures.
Generated result evaluation	Computes metrics such as RMSD, sequence recovery, and (co-)designability.

Usage

1. Using OneCode

You can try intelligent one-click AI4S programming through the OneCode online environment:

Try intelligent one-click AI4S programming

2. Manual Installation and Usage

Hardware Requirements

• Running on a GPU or DCU is recommended.
• CPU can be used for connectivity checks, but it is relatively slow.
• DCU users need to install DTK in advance. DTK 25.04.2 or later is recommended, or the OneScience-recommended version that matches the current cluster.

Software Requirements

DCU users who want to learn more about adaptation details can contact liubiao@sugon.com.

Environment Checks

• NVIDIA GPU:

nvidia-smi

• Hygon DCU:

hy-smi

Environment Preparation

Optional: override the default paths through environment variables:

export LAPROTEINA_ROOT=/path/to/la-proteina
export LAPROTEINA_DATASET_DIR=/path/to/dataset
export LAPROTEINA_CHECKPOINTS_DIR=/path/to/checkpoints_laproteina
export DATA_PATH=/path/to/dataset

Quick Start

1. Install the Runtime Environment

conda create -n onescience311 python=3.11 -y
conda activate onescience311
pip install onescience[bio] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai

If the following code cannot find required libraries at runtime, activate CUDA as shown below.

source ${ROCM_PATH}/cuda/env.sh
export LD_LIBRARY_PATH="$CONDA_PREFIX/lib:$LD_LIBRARY_PATH"
export LD_LIBRARY_PATH="$CONDA_PREFIX/lib/python3.11/site-packages/fastpt/torch/lib:$LD_LIBRARY_PATH"

2. Download the Model Package

# By default, the package is downloaded to the model folder under the current path. To change this, adjust the path after local_dir.
hf download --model OneScience-Sugon/La-Proteina --local-dir ./model
cd model

Training Weights and Datasets

They will be uploaded to Hugging Face soon, and command-line downloads will be supported later.

3. Usage Modes

It is recommended to run the scripts under the laproteina directory so that outputs are managed in one place.

1. Train the Main La-Proteina Model (run_train.sh)

bash scripts/run_train.sh

Common Hydra parameter overrides:

# Single-device debugging
bash scripts/run_train.sh hardware.ngpus_per_node_=1 single=true

# Specify a run name
bash scripts/run_train.sh run_name=my_laproteina_run

# Override dataset or network configuration (for motif training scenarios)
bash scripts/run_train.sh dataset=pdb/pdb_train_motif_aa nn=local_latents_score_nn_160M_motif_idx_aa

Output:

• ./store/<run_name>/: training logs, checkpoints, and Hydra configurations

2. Protein Structure Generation (run_generate.sh)

bash scripts/run_generate.sh

By default, the inference_ucond_tri configuration is used for unconditional generation (LD2 model + AE1).

Switching generation configurations:

# Unconditional generation without triangle attention
bash scripts/run_generate.sh --config_name inference_ucond_notri

# Unconditional generation for long chains (300-800 residues)
bash scripts/run_generate.sh --config_name inference_ucond_notri_long

# Indexed all-atom motif scaffolding
bash scripts/run_generate.sh --config_name inference_motif_idx_aa

# Indexed tip-atom motif scaffolding
bash scripts/run_generate.sh --config_name inference_motif_idx_tip

# Unindexed all-atom motif scaffolding
bash scripts/run_generate.sh --config_name inference_motif_uidx_aa

# Unindexed tip-atom motif scaffolding
bash scripts/run_generate.sh --config_name inference_motif_uidx_tip

Output:

• ./inference/<config_name>/: generated protein structure files and metadata

---

3. Generated Result Evaluation (run_evaluate.sh)

bash scripts/run_evaluate.sh

By default, this evaluates the generated results corresponding to the inference_ucond_tri configuration.

Preparing ProteinMPNN weights:

ProteinMPNN weights must be downloaded before evaluation. They can be downloaded from the Hugging Face community:

hf download --model OneScience-Sugon/ProteinMPNN --local-dir ./weight

Output:

• ./inference/<config_name>/evaluation/: evaluation result files

---

4. Autoencoder Inference (run_ae_infer.sh)

bash scripts/run_ae_infer.sh

This performs encode-decode reconstruction on the PDB dataset and evaluates reconstruction metrics, such as all-atom RMSD and sequence recovery.

The script checks whether DATA_PATH/pdb_train and AE1_ucond_512.ckpt exist.

The script internally calls:

python infer_laproteina_ae.py "$@"

Common parameter overrides:

Environment variable	Default value	Description
LAPROTEINA_ROOT	${ONESCIENCE_DATASETS_DIR}/la-proteina	Root directory for data and weights
LAPROTEINA_CHECKPOINTS_DIR	${LAPROTEINA_ROOT}/checkpoints_laproteina	Autoencoder weight directory
DATA_PATH	${LAPROTEINA_ROOT}/dataset	Dataset directory

Output:

• ./inference_ae/: reconstructed structures and evaluation metrics

Notes

• Make sure the ONESCIENCE_DATASETS_DIR environment variable is correctly set before running the scripts.
• The training script checks whether DATA_PATH/pdb_train and AE1_ucond_512.ckpt exist by default. If either is missing, it will report an error and exit.
• In the current integration, dataset=pdb is available. dataset=genie2 and dataset=pdb_multimer are not packaged in the current repository snapshot, and running them will produce an explicit error.
• The scripts automatically set ROCm/DCU-related LD_LIBRARY_PATH values and can run directly on Hygon DCU platforms. On CUDA platforms, these settings can be ignored or adjusted as needed.
• When generating motif scaffold structures, make sure the LD model is correctly paired with the corresponding AE model. Otherwise, generation may fail due to length or task mismatch.
• Evaluating (co-)designability requires ProteinMPNN weights. Run script_utils/download_pmpnn_weights.sh in advance to download them.
• It is recommended to run all scripts under examples/biosciences/laproteina so that output directories remain unified.
• When overriding configurations through Hydra, you can use +CK_PATH=... to specify the checkpoint root path. If it is not provided, the script automatically sets it to LAPROTEINA_ROOT.

Official OneScience Information

Platform	OneScience main repository	Skills repository
Gitee	https://gitee.com/onescience-ai/onescience	https://gitee.com/onescience-ai/oneskills
GitHub	https://github.com/onescience-ai/OneScience	https://github.com/onescience-ai/oneskills

Citation and License

The example code is licensed under Apache 2.0. La-Proteina model weights are licensed under the NVIDIA Open Model License Agreement, and other materials are licensed under CC-BY 4.0.

If you use La-Proteina in your research, please cite the original paper:

@article{geffner2025laproteina,
  title={La-Proteina: Atomistic Protein Generation via Partially Latent Flow Matching},
  author={Geffner, Tomas and Didi, Kieran and Cao, Zhonglin and Reidenbach, Danny and Zhang, Zuobai and Dallago, Christian and Kucukbenli, Emine and Kreis, Karsten and Vahdat, Arash},
  journal={arXiv preprint arXiv:2507.09466},
  year={2025}
}