Add files using upload-large-folder tool

config/experiment/train_fsdp.yaml

@@ -0,0 +1,54 @@
+# @package _global_
+
+# All parameters below will be merged with parameters from default configurations set above 
+# this allows you to overwrite only specified parameters
+
+# Options for training:
+# - /data: [pdb, pdb_sp, pdb_sp_afesm, pdb_sp_afesme]
+# - /model/architecture: 
+#         [foldingdit_100M, foldingdit_360M, foldingdit_700M, 
+#          foldingdit_1.1B, foldingdit_1.6B, foldingdit_3B]
+# - /trainer: [default, fsdp]
+
+defaults:
+  - override /data: pdb
+  - override /model: simplefold
+  - override /model/architecture: foldingdit_100M
+  - override /model/sampler: euler_maruyama
+  - override /model/processor: protein_processor
+  - override /callbacks: default
+  - override /trainer: fsdp # FSDP training
+  - override /logger: tensorboard
+
+# load_ckpt_path: [PATH_TO_YOUR_CKPT] # e.g. uncomment to load a checkpoint
+seed: 12345
+
+trainer:
+  max_steps: 300000
+  val_check_interval: 10000 # this controls BOTH checkpoint steps and validation
+  accumulate_grad_batches: 1 # always 1 for FSDP
+
+data:
+  batch_size: 8 # should be the same as the number of GPUs
+  num_workers: 16
+
+model:
+  _target_: models.simplefold.simplefold.SimpleFold
+  ema_decay: 0.999
+  clip_grad_norm_val: 2.0
+  esm_model: "esm2_3B"
+
+  scheduler:
+    _target_: onescience.utils.simplefold.lr_scheduler.LinearWarmup
+    _partial_: true
+    min_lr: 1e-6
+    max_lr: ${model.optimizer.lr}
+    warmup_steps: 5000
+
+  architecture:
+    esm_model: ${model.esm_model}
+
+  processor:
+    scale: 16.0
+    ref_scale: 5.0
+    multiplicity: 16 # number of copies per GPU

config/model/architecture/foldingdit_1.1B.yaml

@@ -0,0 +1,99 @@
+_target_: models.simplefold.torch.architecture.FoldingDiT
+
+hidden_size: 1280
+num_heads: 20
+atom_num_heads: 6
+output_channels: 3
+use_atom_mask: False
+use_length_condition: True
+esm_dropout_prob: 0.0
+esm_model: esm2_3B
+
+time_embedder:
+  _target_: models.simplefold.torch.layers.TimestepEmbedder
+  hidden_size: 1280
+aminoacid_pos_embedder:
+  _target_: models.simplefold.torch.pos_embed.AbsolutePositionEncoding
+  in_dim: 1
+  embed_dim: 1280
+  include_input: True
+pos_embedder:
+  _target_: models.simplefold.torch.pos_embed.FourierPositionEncoding
+  in_dim: 3
+  include_input: True
+  min_freq_log2: 0
+  max_freq_log2: 12
+  num_freqs: 128
+  log_sampling: True
+
+trunk:
+  _target_: models.simplefold.torch.blocks.HomogenTrunk
+  depth: 36
+  block: 
+    _target_: models.simplefold.torch.blocks.DiTBlock
+    _partial_: True # because in the for loop we create a new module
+    hidden_size: 1280
+    mlp_ratio: 4.0
+    use_swiglu: True # SwiGLU FFN
+    self_attention_layer:
+      _target_: models.simplefold.torch.layers.EfficientSelfAttentionLayer
+      _partial_: True
+      hidden_size: 1280
+      num_heads: 20
+      qk_norm: True
+      pos_embedder:
+        _target_: models.simplefold.torch.pos_embed.AxialRotaryPositionEncoding
+        in_dim: 4
+        embed_dim: 1280
+        num_heads: 20
+        base: 100.0
+
+atom_hidden_size_enc: 384
+atom_n_queries_enc: 32
+atom_n_keys_enc: 128
+atom_encoder_transformer:
+  _target_: models.simplefold.torch.blocks.HomogenTrunk
+  depth: 2
+  block: 
+    _target_: models.simplefold.torch.blocks.DiTBlock
+    _partial_: True # because in the for loop we create a new module
+    hidden_size: 384
+    mlp_ratio: 4.0
+    use_swiglu: True # SwiGLU FFN
+    self_attention_layer:
+      _target_: models.simplefold.torch.layers.EfficientSelfAttentionLayer
+      _partial_: True
+      hidden_size: 384
+      num_heads: 6
+      qk_norm: True
+      pos_embedder:
+        _target_: models.simplefold.torch.pos_embed.AxialRotaryPositionEncoding
+        in_dim: 4
+        embed_dim: 384
+        num_heads: 6
+        base: 100.0
+
+atom_hidden_size_dec: 384
+atom_n_queries_dec: 32
+atom_n_keys_dec: 128
+atom_decoder_transformer:
+  _target_: models.simplefold.torch.blocks.HomogenTrunk
+  depth: 2
+  block: 
+    _target_: models.simplefold.torch.blocks.DiTBlock
+    _partial_: True # because in the for loop we create a new module
+    hidden_size: 384
+    mlp_ratio: 4.0
+    use_swiglu: True # SwiGLU FFN
+    self_attention_layer:
+      _target_: models.simplefold.torch.layers.EfficientSelfAttentionLayer
+      _partial_: True
+      hidden_size: 384
+      num_heads: 6
+      qk_norm: True
+      pos_embedder:
+        _target_: models.simplefold.torch.pos_embed.AxialRotaryPositionEncoding
+        in_dim: 4
+        embed_dim: 384
+        num_heads: 6
+        base: 100.0

config/model/architecture/foldingdit_100M.yaml

@@ -0,0 +1,101 @@
+_target_: models.simplefold.torch.architecture.FoldingDiT
+
+hidden_size: 768
+num_heads: 12
+atom_num_heads: 4
+output_channels: 3
+use_atom_mask: False
+use_length_condition: True
+esm_dropout_prob: 0.0
+esm_model: esm2_3B
+
+time_embedder:
+  _target_: models.simplefold.torch.layers.TimestepEmbedder
+  hidden_size: 768
+
+aminoacid_pos_embedder:
+  _target_: models.simplefold.torch.pos_embed.AbsolutePositionEncoding
+  in_dim: 1
+  embed_dim: 768
+  include_input: True
+
+pos_embedder:
+  _target_: models.simplefold.torch.pos_embed.FourierPositionEncoding
+  in_dim: 3
+  include_input: True
+  min_freq_log2: 0
+  max_freq_log2: 12
+  num_freqs: 128
+  log_sampling: True
+
+trunk:
+  _target_: models.simplefold.torch.blocks.HomogenTrunk
+  depth: 8
+  block: 
+    _target_: models.simplefold.torch.blocks.DiTBlock
+    _partial_: True # because in the for loop we create a new module
+    hidden_size: 768
+    mlp_ratio: 4.0
+    use_swiglu: True # SwiGLU FFN
+    self_attention_layer:
+      _target_: models.simplefold.torch.layers.EfficientSelfAttentionLayer
+      _partial_: True
+      hidden_size: 768
+      num_heads: 12
+      qk_norm: True
+      pos_embedder:
+        _target_: models.simplefold.torch.pos_embed.AxialRotaryPositionEncoding
+        in_dim: 4
+        embed_dim: 768
+        num_heads: 12
+        base: 100.0
+
+atom_hidden_size_enc: 256
+atom_n_queries_enc: 32
+atom_n_keys_enc: 128
+atom_encoder_transformer:
+  _target_: models.simplefold.torch.blocks.HomogenTrunk
+  depth: 1
+  block: 
+    _target_: models.simplefold.torch.blocks.DiTBlock
+    _partial_: True # because in the for loop we create a new module
+    hidden_size: 256
+    mlp_ratio: 4.0
+    use_swiglu: True # SwiGLU FFN
+    self_attention_layer:
+      _target_: models.simplefold.torch.layers.EfficientSelfAttentionLayer
+      _partial_: True
+      hidden_size: 256
+      num_heads: 4
+      qk_norm: True
+      pos_embedder:
+        _target_: models.simplefold.torch.pos_embed.AxialRotaryPositionEncoding
+        in_dim: 4
+        embed_dim: 256
+        num_heads: 4
+        base: 100.0
+
+atom_hidden_size_dec: 256
+atom_n_queries_dec: 32
+atom_n_keys_dec: 128
+atom_decoder_transformer:
+  _target_: models.simplefold.torch.blocks.HomogenTrunk
+  depth: 1
+  block: 
+    _target_: models.simplefold.torch.blocks.DiTBlock
+    _partial_: True # because in the for loop we create a new module
+    hidden_size: 256
+    mlp_ratio: 4.0
+    use_swiglu: True # SwiGLU FFN
+    self_attention_layer:
+      _target_: models.simplefold.torch.layers.EfficientSelfAttentionLayer
+      _partial_: True
+      hidden_size: 256
+      num_heads: 4
+      qk_norm: True
+      pos_embedder:
+        _target_: models.simplefold.torch.pos_embed.AxialRotaryPositionEncoding
+        in_dim: 4
+        embed_dim: 256
+        num_heads: 4
+        base: 100.0

config/model/architecture/foldingdit_360M.yaml

@@ -0,0 +1,101 @@
+_target_: models.simplefold.torch.architecture.FoldingDiT
+
+hidden_size: 1024
+num_heads: 16
+atom_num_heads: 4
+output_channels: 3
+use_atom_mask: False
+use_length_condition: True
+esm_dropout_prob: 0.0
+esm_model: esm2_3B
+
+time_embedder:
+  _target_: models.simplefold.torch.layers.TimestepEmbedder
+  hidden_size: 1024
+
+aminoacid_pos_embedder:
+  _target_: models.simplefold.torch.pos_embed.AbsolutePositionEncoding
+  in_dim: 1
+  embed_dim: 1024
+  include_input: True
+
+pos_embedder:
+  _target_: models.simplefold.torch.pos_embed.FourierPositionEncoding
+  in_dim: 3
+  include_input: True
+  min_freq_log2: 0
+  max_freq_log2: 12
+  num_freqs: 128
+  log_sampling: True
+
+trunk:
+  _target_: models.simplefold.torch.blocks.HomogenTrunk
+  depth: 18
+  block: 
+    _target_: models.simplefold.torch.blocks.DiTBlock
+    _partial_: True # because in the for loop we create a new module
+    hidden_size: 1024
+    mlp_ratio: 4.0
+    use_swiglu: True # SwiGLU FFN
+    self_attention_layer:
+      _target_: models.simplefold.torch.layers.EfficientSelfAttentionLayer
+      _partial_: True
+      hidden_size: 1024
+      num_heads: 16
+      qk_norm: True
+      pos_embedder:
+        _target_: models.simplefold.torch.pos_embed.AxialRotaryPositionEncoding
+        in_dim: 4
+        embed_dim: 1024
+        num_heads: 16
+        base: 100.0
+
+atom_hidden_size_enc: 256
+atom_n_queries_enc: 32
+atom_n_keys_enc: 128
+atom_encoder_transformer:
+  _target_: models.simplefold.torch.blocks.HomogenTrunk
+  depth: 2
+  block: 
+    _target_: models.simplefold.torch.blocks.DiTBlock
+    _partial_: True # because in the for loop we create a new module
+    hidden_size: 256
+    mlp_ratio: 4.0
+    use_swiglu: True # SwiGLU FFN
+    self_attention_layer:
+      _target_: models.simplefold.torch.layers.EfficientSelfAttentionLayer
+      _partial_: True
+      hidden_size: 256
+      num_heads: 4
+      qk_norm: True
+      pos_embedder:
+        _target_: models.simplefold.torch.pos_embed.AxialRotaryPositionEncoding
+        in_dim: 4
+        embed_dim: 256
+        num_heads: 4
+        base: 100.0
+
+atom_hidden_size_dec: 256
+atom_n_queries_dec: 32
+atom_n_keys_dec: 128
+atom_decoder_transformer:
+  _target_: models.simplefold.torch.blocks.HomogenTrunk
+  depth: 2
+  block: 
+    _target_: models.simplefold.torch.blocks.DiTBlock
+    _partial_: True # because in the for loop we create a new module
+    hidden_size: 256
+    mlp_ratio: 4.0
+    use_swiglu: True # SwiGLU FFN
+    self_attention_layer:
+      _target_: models.simplefold.torch.layers.EfficientSelfAttentionLayer
+      _partial_: True
+      hidden_size: 256
+      num_heads: 4
+      qk_norm: True
+      pos_embedder:
+        _target_: models.simplefold.torch.pos_embed.AxialRotaryPositionEncoding
+        in_dim: 4
+        embed_dim: 256
+        num_heads: 4
+        base: 100.0

config/model/architecture/foldingdit_3B.yaml

@@ -0,0 +1,99 @@
+_target_: models.simplefold.torch.architecture.FoldingDiT
+
+hidden_size: 2048
+num_heads: 32
+atom_num_heads: 10
+output_channels: 3
+use_atom_mask: False
+use_length_condition: True
+esm_dropout_prob: 0.0
+esm_model: esm2_3B
+
+time_embedder:
+  _target_: models.simplefold.torch.layers.TimestepEmbedder
+  hidden_size: 2048
+aminoacid_pos_embedder:
+  _target_: models.simplefold.torch.pos_embed.AbsolutePositionEncoding
+  in_dim: 1
+  embed_dim: 2048
+  include_input: True
+pos_embedder:
+  _target_: models.simplefold.torch.pos_embed.FourierPositionEncoding
+  in_dim: 3
+  include_input: True
+  min_freq_log2: 0
+  max_freq_log2: 12
+  num_freqs: 128
+  log_sampling: True
+
+trunk:
+  _target_: models.simplefold.torch.blocks.HomogenTrunk
+  depth: 36
+  block: 
+    _target_: models.simplefold.torch.blocks.DiTBlock
+    _partial_: True # because in the for loop we create a new module
+    hidden_size: 2048
+    mlp_ratio: 4.0
+    use_swiglu: True # SwiGLU FFN
+    self_attention_layer:
+      _target_: models.simplefold.torch.layers.EfficientSelfAttentionLayer
+      _partial_: True
+      hidden_size: 2048
+      num_heads: 32
+      qk_norm: True
+      pos_embedder:
+        _target_: models.simplefold.torch.pos_embed.AxialRotaryPositionEncoding
+        in_dim: 4
+        embed_dim: 2048
+        num_heads: 32
+        base: 100.0
+
+atom_hidden_size_enc: 640
+atom_n_queries_enc: 32
+atom_n_keys_enc: 128
+atom_encoder_transformer:
+  _target_: models.simplefold.torch.blocks.HomogenTrunk
+  depth: 4
+  block: 
+    _target_: models.simplefold.torch.blocks.DiTBlock
+    _partial_: True # because in the for loop we create a new module
+    hidden_size: 640
+    mlp_ratio: 4.0
+    use_swiglu: True # SwiGLU FFN
+    self_attention_layer:
+      _target_: models.simplefold.torch.layers.EfficientSelfAttentionLayer
+      _partial_: True
+      hidden_size: 640
+      num_heads: 10
+      qk_norm: True
+      pos_embedder:
+        _target_: models.simplefold.torch.pos_embed.AxialRotaryPositionEncoding
+        in_dim: 4
+        embed_dim: 640
+        num_heads: 10
+        base: 100.0
+
+atom_hidden_size_dec: 640
+atom_n_queries_dec: 32
+atom_n_keys_dec: 128
+atom_decoder_transformer:
+  _target_: models.simplefold.torch.blocks.HomogenTrunk
+  depth: 4
+  block: 
+    _target_: models.simplefold.torch.blocks.DiTBlock
+    _partial_: True # because in the for loop we create a new module
+    hidden_size: 640
+    mlp_ratio: 4.0
+    use_swiglu: True # SwiGLU FFN
+    self_attention_layer:
+      _target_: models.simplefold.torch.layers.EfficientSelfAttentionLayer
+      _partial_: True
+      hidden_size: 640
+      num_heads: 10
+      qk_norm: True
+      pos_embedder:
+        _target_: models.simplefold.torch.pos_embed.AxialRotaryPositionEncoding
+        in_dim: 4
+        embed_dim: 640
+        num_heads: 10
+        base: 100.0

config/model/architecture/foldingdit_700M.yaml

@@ -0,0 +1,101 @@
+_target_: models.simplefold.torch.architecture.FoldingDiT
+
+hidden_size: 1152
+num_heads: 16
+atom_num_heads: 4
+output_channels: 3
+use_atom_mask: False
+use_length_condition: True
+esm_dropout_prob: 0.0
+esm_model: esm2_3B
+
+time_embedder:
+  _target_: models.simplefold.torch.layers.TimestepEmbedder
+  hidden_size: 1152
+
+aminoacid_pos_embedder:
+  _target_: models.simplefold.torch.pos_embed.AbsolutePositionEncoding
+  in_dim: 1
+  embed_dim: 1152
+  include_input: True
+
+pos_embedder:
+  _target_: models.simplefold.torch.pos_embed.FourierPositionEncoding
+  in_dim: 3
+  include_input: True
+  min_freq_log2: 0
+  max_freq_log2: 12
+  num_freqs: 128
+  log_sampling: True
+
+trunk:
+  _target_: models.simplefold.torch.blocks.HomogenTrunk
+  depth: 28
+  block: 
+    _target_: models.simplefold.torch.blocks.DiTBlock
+    _partial_: True # because in the for loop we create a new module
+    hidden_size: 1152
+    mlp_ratio: 4.0
+    use_swiglu: True # SwiGLU FFN
+    self_attention_layer:
+      _target_: models.simplefold.torch.layers.EfficientSelfAttentionLayer
+      _partial_: True
+      hidden_size: 1152
+      num_heads: 16
+      qk_norm: True
+      pos_embedder:
+        _target_: models.simplefold.torch.pos_embed.AxialRotaryPositionEncoding
+        in_dim: 4
+        embed_dim: 1152
+        num_heads: 16
+        base: 100.0
+
+atom_hidden_size_enc: 256
+atom_n_queries_enc: 32
+atom_n_keys_enc: 128
+atom_encoder_transformer:
+  _target_: models.simplefold.torch.blocks.HomogenTrunk
+  depth: 2
+  block: 
+    _target_: models.simplefold.torch.blocks.DiTBlock
+    _partial_: True # because in the for loop we create a new module
+    hidden_size: 256
+    mlp_ratio: 4.0
+    use_swiglu: True # SwiGLU FFN
+    self_attention_layer:
+      _target_: models.simplefold.torch.layers.EfficientSelfAttentionLayer
+      _partial_: True
+      hidden_size: 256
+      num_heads: 4
+      qk_norm: True
+      pos_embedder:
+        _target_: models.simplefold.torch.pos_embed.AxialRotaryPositionEncoding
+        in_dim: 4
+        embed_dim: 256
+        num_heads: 4
+        base: 100.0
+
+atom_hidden_size_dec: 256
+atom_n_queries_dec: 32
+atom_n_keys_dec: 128
+atom_decoder_transformer:
+  _target_: models.simplefold.torch.blocks.HomogenTrunk
+  depth: 2
+  block: 
+    _target_: models.simplefold.torch.blocks.DiTBlock
+    _partial_: True # because in the for loop we create a new module
+    hidden_size: 256
+    mlp_ratio: 4.0
+    use_swiglu: True # SwiGLU FFN
+    self_attention_layer:
+      _target_: models.simplefold.torch.layers.EfficientSelfAttentionLayer
+      _partial_: True
+      hidden_size: 256
+      num_heads: 4
+      qk_norm: True
+      pos_embedder:
+        _target_: models.simplefold.torch.pos_embed.AxialRotaryPositionEncoding
+        in_dim: 4
+        embed_dim: 256
+        num_heads: 4
+        base: 100.0

config/model/sampler/euler_maruyama.yaml

@@ -0,0 +1,6 @@
+_target_: models.simplefold.torch.sampler.EMSampler
+num_timesteps: 500
+t_start: 1e-4
+tau: 0.3
+log_timesteps: True
+w_cutoff: 0.99

config/paths/default.yaml

@@ -0,0 +1,20 @@
+# path to root directory
+# this requires PROJECT_ROOT environment variable to exist
+# you can replace it with "." if you want the root to be the current working directory
+root_dir: "./"
+
+tmp_dir: ${paths.root_dir}tmp/
+
+# path to data directory
+data_dir: ${paths.root_dir}data/
+
+# path to logging directory
+log_dir: ${paths.root_dir}logs/
+
+# artifacts directory
+output_dir: ${paths.root_dir}artifacts/
+
+sample_dir: ${.output_dir}/samples/
+
+# path to working directory
+work_dir: ${hydra:runtime.cwd}

examples/minimal.fasta

@@ -0,0 +1,2 @@
+>simplefold_minimal_example
+GASKLRAVLEKLKLSRDDISTAAGMVKGVVDHLLLRLKCDSAFRGVGLLNTGSYYEHVKISAPNEFDVMFKLEVPRIQLEEYSNTRAYYFVKFKRNPKENPLSQFLEGEILSASKMLSKFRKIIKEEINDDTDVIMKRKRGGSPAVTLLISEKISVDITLALESKSSWPASTQEGLRIQNWLSAKVRKQLRLKPFYLVPKHAEETWRLSFSHIEKEILNNHGKSKTCCENKEEKCCRKDCLKLMKYLLEQLKERFKDKKHLDKFSSYHVKTAFFHVCTQNPQDSQWDRKDLGLCFDNCVTYFLQCLRTEKLENYFIPEFNLFSSNLIDKRSKEFLTKQIEYERNNEFPVFD

models/__init__.py

@@ -0,0 +1,2 @@
+"""Bundled model modules."""
+

models/simplefold/__init__.py

@@ -0,0 +1,4 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+#

models/simplefold/flow.py

@@ -0,0 +1,101 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+#
+
+def right_pad_dims_to(x, t):
+    padding_dims = x.ndim - t.ndim
+    if padding_dims <= 0:
+        return t
+    return t.reshape(*t.shape, *((1,) * padding_dims))
+
+
+class BasePath:
+    """base class for flow matching path"""
+
+    def __init__(self):
+        return
+
+    def compute_alpha_t(self, t):
+        """Compute the data coefficient along the path"""
+        return None, None
+
+    def compute_sigma_t(self, t):
+        """Compute the noise coefficient along the path"""
+        return None, None
+
+    def compute_d_alpha_alpha_ratio_t(self, t):
+        """Compute the ratio between d_alpha and alpha"""
+        alpha_t, d_alpha_t = self.compute_alpha_t(t)
+        return d_alpha_t / alpha_t
+
+    def compute_mu_t(self, t, x0, x1):
+        """Compute the mean of time-dependent density p_t"""
+        alpha_t, _ = self.compute_alpha_t(t)
+        sigma_t, _ = self.compute_sigma_t(t)
+        return alpha_t * x1 + sigma_t * x0
+
+    def compute_xt(self, t, x0, x1):
+        """Sample xt from time-dependent density p_t; rng is required"""
+        xt = self.compute_mu_t(t, x0, x1)
+        return xt
+
+    def compute_ut(self, t, x0, x1):
+        """Compute the vector field corresponding to p_t"""
+        _, d_alpha_t = self.compute_alpha_t(t)
+        _, d_sigma_t = self.compute_sigma_t(t)
+        return d_alpha_t * x1 + d_sigma_t * x0
+
+    def interpolant(self, t, x0, x1):
+        t = right_pad_dims_to(x0, t)
+        xt = self.compute_xt(t, x0, x1)
+        ut = self.compute_ut(t, x0, x1)
+        return t, xt, ut
+
+    def compute_drift(self, x, t):
+        """We always output sde according to score parametrization; """
+        t = right_pad_dims_to(x, t)
+        alpha_ratio = self.compute_d_alpha_alpha_ratio_t(t)
+        sigma_t, d_sigma_t = self.compute_sigma_t(t)
+        drift_mean = alpha_ratio * x
+        drift_var = alpha_ratio * (sigma_t ** 2) - sigma_t * d_sigma_t
+        return -drift_mean, drift_var
+
+    def compute_score_from_velocity(self, v_t, y_t, t):
+        t = right_pad_dims_to(y_t, t)
+        alpha_t, d_alpha_t = self.compute_alpha_t(t)
+        sigma_t, d_sigma_t = self.compute_sigma_t(t)
+        mean = y_t
+        reverse_alpha_ratio = alpha_t / d_alpha_t
+        var = sigma_t**2 - reverse_alpha_ratio * d_sigma_t * sigma_t
+        score = (reverse_alpha_ratio * v_t - mean) / var
+        return score
+
+    def compute_velocity_from_score(self, s_t, y_t, t):
+        t = right_pad_dims_to(y_t, t)
+        drift_mean, drift_var = self.compute_drift(y_t, t)
+        velocity = -drift_mean + drift_var * s_t
+        return velocity
+
+
+class LinearPath(BasePath):
+    """
+    Linear flow process:
+    x0: noise, x1: data
+    In inference, we sample data from 0 -> 1
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def compute_alpha_t(self, t):
+        """Compute the data coefficient along the path"""
+        return t, 1
+
+    def compute_sigma_t(self, t):
+        """Compute the noise coefficient along the path"""
+        return 1 - t, -1
+
+    def compute_d_alpha_alpha_ratio_t(self, t):
+        """Compute the ratio between d_alpha and alpha"""
+        return 1 / t

models/simplefold/mlx/__init__.py

@@ -0,0 +1,4 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+#

models/simplefold/mlx/architecture.py

@@ -0,0 +1,333 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+#
+
+import math
+import mlx.nn as nn
+import mlx.core as mx
+from models.simplefold.mlx.layers import FinalLayer, ConditionEmbedder
+from onescience.utils.simplefold.esm_utils import esm_model_dict
+
+
+# MLX does not have a native one_hot implementation
+def one_hot(indices, num_classes, dtype=None):
+    """
+    MLX version of torch.one_hot.
+
+    Args:
+        indices: integer MLX array of any shape, containing class indices in [0, num_classes).
+        num_classes: number of classes for the one-hot dimension.
+        dtype: output data type (defaults to float32).
+
+    Returns:
+        MLX array of shape indices.shape + (num_classes,) and given dtype.
+    """
+    # Default to float32 if no dtype is given
+    if dtype is None:
+        dtype = mx.float32
+
+    classes = mx.arange(num_classes, dtype=indices.dtype)
+
+    # Broadcast-compare: result has shape indices.shape + (num_classes,)
+    # For each position, only the matched class index gives True
+    mask = indices[..., mx.newaxis] == classes
+
+    # Cast boolean mask to desired dtype
+    return mask.astype(dtype)
+
+
+class FoldingDiT(nn.Module):
+    def __init__(
+        self,
+        trunk,
+        time_embedder,
+        aminoacid_pos_embedder,
+        pos_embedder,
+        atom_encoder_transformer,
+        atom_decoder_transformer,
+        hidden_size=1152,
+        num_heads=16,
+        atom_num_heads=4,
+        output_channels=3,
+        atom_hidden_size_enc=256,
+        atom_hidden_size_dec=256,
+        atom_n_queries_enc=32,
+        atom_n_keys_enc=128,
+        atom_n_queries_dec=32,
+        atom_n_keys_dec=128,
+        esm_model="esm2_3B",
+        esm_dropout_prob=0.0,
+        use_atom_mask=False,
+        use_length_condition=True,
+    ):
+        super().__init__()
+        self.pos_embedder = pos_embedder
+        pos_embed_channels = pos_embedder.embed_dim
+        self.aminoacid_pos_embedder = aminoacid_pos_embedder
+        aminoacid_pos_embed_channels = aminoacid_pos_embedder.embed_dim
+
+        self.time_embedder = time_embedder
+
+        self.atom_encoder_transformer = atom_encoder_transformer
+        self.atom_decoder_transformer = atom_decoder_transformer
+
+        self.trunk = trunk
+
+        self.hidden_size = hidden_size
+        self.output_channels = output_channels
+        self.num_heads = num_heads
+        self.atom_num_heads = atom_num_heads
+        self.use_atom_mask = use_atom_mask
+        self.esm_dropout_prob = esm_dropout_prob
+        self.use_length_condition = use_length_condition
+
+        esm_s_dim = esm_model_dict[esm_model]["esm_s_dim"]
+        esm_num_layers = esm_model_dict[esm_model]["esm_num_layers"]
+
+        self.atom_hidden_size_enc = atom_hidden_size_enc
+        self.atom_hidden_size_dec = atom_hidden_size_dec
+        self.atom_n_queries_enc = atom_n_queries_enc
+        self.atom_n_keys_enc = atom_n_keys_enc
+        self.atom_n_queries_dec = atom_n_queries_dec
+        self.atom_n_keys_dec = atom_n_keys_dec
+
+        atom_feat_dim = pos_embed_channels + aminoacid_pos_embed_channels + 427
+        self.atom_feat_proj = nn.Sequential(
+            nn.Linear(atom_feat_dim, hidden_size),
+            nn.LayerNorm(hidden_size),
+            nn.SiLU(),
+        )
+        self.atom_pos_proj = nn.Linear(pos_embed_channels, hidden_size, bias=False)
+
+        if self.use_length_condition:
+            self.length_embedder = nn.Sequential(
+                nn.Linear(1, hidden_size, bias=False),
+                nn.LayerNorm(hidden_size),
+            )
+
+        self.atom_in_proj = nn.Linear(hidden_size * 2, hidden_size, bias=False)
+
+        self.esm_s_combine = mx.zeros(esm_num_layers)
+        self.esm_s_proj = ConditionEmbedder(
+            input_dim=esm_s_dim,
+            hidden_size=hidden_size,
+            dropout_prob=0,
+        )
+
+        latent_cat_dim = hidden_size * 2
+        self.esm_cat_proj = nn.Linear(latent_cat_dim, hidden_size)
+
+        self.context2atom_proj = nn.Sequential(
+            nn.Linear(hidden_size, self.atom_hidden_size_enc),
+            nn.LayerNorm(self.atom_hidden_size_enc),
+        )
+
+        self.atom_enc_cond_proj = nn.Sequential(
+            nn.Linear(hidden_size, self.atom_hidden_size_enc),
+            nn.LayerNorm(self.atom_hidden_size_enc),
+        )
+
+        self.atom2latent_proj = nn.Sequential(
+            nn.Linear(self.atom_hidden_size_enc, hidden_size),
+            nn.LayerNorm(hidden_size),
+        )
+
+        self.atom_dec_cond_proj = nn.Sequential(
+            nn.Linear(hidden_size, self.atom_hidden_size_dec),
+            nn.LayerNorm(self.atom_hidden_size_dec),
+        )
+
+        self.latent2atom_proj = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size),
+            nn.SiLU(),
+            nn.LayerNorm(hidden_size),
+            nn.Linear(hidden_size, self.atom_hidden_size_dec),
+        )
+
+        self.final_layer = FinalLayer(
+            self.atom_hidden_size_dec, output_channels, c_dim=hidden_size
+        )
+
+    def create_local_attn_bias(
+        self,
+        n: int,
+        n_queries: int,
+        n_keys: int,
+        inf: float = 1e10,
+    ):
+        """Create local attention bias based on query window n_queries and kv window n_keys.
+
+        Args:
+            n (int): the length of quiries
+            n_queries (int): window size of quiries
+            n_keys (int): window size of keys/values
+            inf (float, optional): the inf to mask attention. Defaults to 1e10.
+            device (torch.device, optional): cuda|cpu|None. Defaults to None.
+
+        Returns:
+            torch.Tensor: the diagonal-like global attention bias
+        """
+        n_trunks = int(math.ceil(n / n_queries))
+        padded_n = n_trunks * n_queries
+        attn_mask = mx.zeros((padded_n, padded_n))
+        for block_index in range(0, n_trunks):
+            i = block_index * n_queries
+            j1 = max(0, n_queries * block_index - (n_keys - n_queries) // 2)
+            j2 = n_queries * block_index + (n_queries + n_keys) // 2
+            attn_mask[i : i + n_queries, j1:j2] = 1.0
+        attn_bias = (1 - attn_mask) * -inf
+        return attn_bias[:n, :n]
+
+    def create_atom_attn_mask(
+        self, feats, natoms, atom_n_queries=None, atom_n_keys=None, inf: float = 1e10
+    ):
+
+        if atom_n_queries is not None and atom_n_keys is not None:
+            atom_attn_mask = self.create_local_attn_bias(
+                n=natoms, n_queries=atom_n_queries, n_keys=atom_n_keys, inf=inf
+            )
+        else:
+            atom_attn_mask = None
+
+        return atom_attn_mask
+
+    def __call__(self, noised_pos, t, feats, self_cond=None):
+
+        B, N, _ = feats["ref_pos"].shape
+        M = feats["mol_type"].shape[1]
+        atom_to_token = feats["atom_to_token"].astype(mx.float32)
+        atom_to_token_idx = feats["atom_to_token_idx"]
+        ref_space_uid = feats["ref_space_uid"]
+
+        # create atom attention masks
+        atom_attn_mask_enc = self.create_atom_attn_mask(
+            feats,
+            natoms=N,
+            atom_n_queries=self.atom_n_queries_enc,
+            atom_n_keys=self.atom_n_keys_enc,
+        )
+        atom_attn_mask_dec = self.create_atom_attn_mask(
+            feats,
+            natoms=N,
+            atom_n_queries=self.atom_n_queries_dec,
+            atom_n_keys=self.atom_n_keys_dec,
+        )
+
+        # create condition embeddings for AdaLN
+        c_emb = self.time_embedder(t)  # (B, D)
+        if self.use_length_condition:
+            length = feats["max_num_tokens"].astype(mx.float32)[..., None]
+            c_emb = c_emb + self.length_embedder(mx.log(length))
+
+        mol_type = feats["mol_type"]
+        mol_type = one_hot(mol_type, num_classes=4).astype(mx.float32)  # [B, M, 4]
+        res_type = feats["res_type"].astype(mx.float32)  # [B, M, 33]
+        pocket_feature = feats["pocket_feature"].astype(mx.float32)  # [B, M, 4]
+        res_feat = mx.concatenate(
+            [mol_type, res_type, pocket_feature], axis=-1
+        )  # [B, M, 41]
+        atom_feat_from_res = mx.matmul(atom_to_token, res_feat)  # [B, N, 41]
+        atom_res_pos = self.aminoacid_pos_embedder(
+            pos=atom_to_token_idx[..., None].astype(mx.float32)
+        )
+        ref_pos_emb = self.pos_embedder(pos=feats["ref_pos"])
+        atom_feat = mx.concatenate(
+            [
+                ref_pos_emb,  # (B, N, PD1)
+                atom_feat_from_res,  # (B, N, 41)
+                atom_res_pos,  # (B, N, PD2)
+                feats["ref_charge"][..., None],  # (B, N, 1)
+                feats["atom_pad_mask"][..., None],  # (B, N, 1)
+                feats["ref_element"],  # (B, N, 128)
+                feats["ref_atom_name_chars"].reshape(B, N, 4 * 64),  # (B, N, 256)
+            ],
+            axis=-1,
+        )  # (B, N, PD1+PD2+427)
+        atom_feat = self.atom_feat_proj(atom_feat)  # (B, N, D)
+
+        atom_coord = self.pos_embedder(pos=noised_pos)  # (B, N, PD1)
+        atom_coord = self.atom_pos_proj(atom_coord)  # (B, N, D)
+
+        atom_in = mx.concatenate([atom_feat, atom_coord], axis=-1)
+        atom_in = self.atom_in_proj(atom_in)  # (B, N, D)
+
+        # position embeddings for Axial RoPE
+        atom_pe_pos = mx.concatenate(
+            [
+                ref_space_uid[..., None].astype(mx.float32),  # (B, N, 1)
+                feats["ref_pos"],  # (B, N, 3)
+            ],
+            axis=-1,
+        )  # (B, N, 4)
+
+        token_pe_pos = mx.concatenate(
+            [
+                feats["residue_index"][..., None].astype(mx.float32),  # (B, M, 1)
+                feats["entity_id"][..., None].astype(mx.float32),  # (B, M, 1)
+                feats["asym_id"][..., None].astype(mx.float32),  # (B, M, 1)
+                feats["sym_id"][..., None].astype(mx.float32),  # (B, M, 1)
+            ],
+            axis=-1,
+        )  # (B, M, 4)
+
+        atom_c_emb_enc = self.atom_enc_cond_proj(c_emb)
+        atom_latent = self.context2atom_proj(atom_in)
+        atom_latent = self.atom_encoder_transformer(
+            latents=atom_latent,
+            c=atom_c_emb_enc,
+            attention_mask=atom_attn_mask_enc,
+            pos=atom_pe_pos,
+        )
+
+
+        atom_latent = self.atom2latent_proj(atom_latent)
+
+        # grouping: aggregate atom tokens to residue tokens
+        atom_to_token_mean = atom_to_token / (
+            atom_to_token.sum(axis=1, keepdims=True) + 1e-6
+        )
+        latent = mx.matmul(atom_to_token_mean.swapaxes(axis1=1, axis2=2), atom_latent)
+        assert latent.shape[1] == M
+
+        esm_s = (
+            mx.softmax(self.esm_s_combine, axis=0)[None, ...] @ feats["esm_s"]
+        ).squeeze(axis=2)
+
+        # MLX is only interended for inference, we do not drop any ids
+        esm_emb = self.esm_s_proj(esm_s, train=False)
+        assert esm_emb.shape[1] == latent.shape[1]
+
+        latent = self.esm_cat_proj(mx.concatenate([latent, esm_emb], axis=-1))
+
+        # residue trunk
+        latent = self.trunk(
+            latents=latent,
+            c=c_emb,
+            attention_mask=None,
+            pos=token_pe_pos,
+        )
+
+        # ungrouping: broadcast residue tokens to atom tokens
+        output = mx.matmul(atom_to_token, latent)
+        assert output.shape[1] == N
+
+        # add skip connection
+        output = output + atom_latent
+        output = self.latent2atom_proj(output)
+
+        # atom decoder
+        atom_c_emb_dec = self.atom_dec_cond_proj(c_emb)
+        output = self.atom_decoder_transformer(
+            latents=output,
+            c=atom_c_emb_dec,
+            attention_mask=atom_attn_mask_dec,
+            pos=atom_pe_pos,
+        )
+
+        output = self.final_layer(output, c=c_emb)
+
+        return {
+            "predict_velocity": output,
+            "latent": latent,
+        }

models/simplefold/mlx/blocks.py

@@ -0,0 +1,100 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+#
+
+import mlx.nn as nn
+import mlx.core as mx
+
+from .simplefold.mlx.layers import modulate, SwiGLUFeedForward
+
+
+class DiTBlock(nn.Module):
+    """
+    A DiT block with adaptive layer norm zero (adaLN-Zero) conditioning.
+    """
+
+    def __init__(
+        self,
+        self_attention_layer,
+        hidden_size,
+        mlp_ratio=4.0,
+        use_swiglu=True,
+    ):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_size, affine=False, eps=1e-6)
+        self.attn = self_attention_layer()
+        self.norm2 = nn.LayerNorm(hidden_size, affine=False, eps=1e-6)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+
+        assert use_swiglu, "Need use_swiglu=True for MLX"
+        self.mlp = SwiGLUFeedForward(hidden_size, mlp_hidden_dim)
+
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True)
+        )
+
+    def __call__(
+        self,
+        latents,
+        c,
+        **kwargs,
+    ):
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+            self.adaLN_modulation(c).split(6, axis=1)
+        )
+
+        _latents, _ = self.attn(
+            modulate(self.norm1(latents), shift_msa, scale_msa), **kwargs
+        )
+        latents = latents + mx.expand_dims(gate_msa, axis=1) * _latents
+
+        latents = latents + mx.expand_dims(gate_mlp, axis=1) * self.mlp(
+            modulate(self.norm2(latents), shift_mlp, scale_mlp)
+        )
+        return latents
+
+
+class TransformerBlock(nn.Module):
+    """
+    A DiT block with adaptive layer norm zero (adaLN-Zero) conditioning.
+    """
+
+    def __init__(
+        self,
+        self_attention_layer,
+        hidden_size,
+        mlp_ratio=4.0,
+        use_swiglu=False,
+    ):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_size, affine=False, eps=1e-6)
+        self.attn = self_attention_layer()
+        self.norm2 = nn.LayerNorm(hidden_size, affine=False, eps=1e-6)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+
+        assert use_swiglu, "Need use_swiglu=True for MLX"
+        self.mlp = SwiGLUFeedForward(hidden_size, mlp_hidden_dim)
+
+    def __call__(
+        self,
+        latents,
+        **kwargs,
+    ):
+        _latents, _ = self.attn(self.norm1(latents), **kwargs)
+        latents = latents + _latents
+        latents = latents + self.mlp(self.norm2(latents))
+        return latents
+
+
+# Homogen trunk, same block applied iteratively
+class HomogenTrunk(nn.Module):
+    def __init__(self, block, depth):
+        super().__init__()
+        self.blocks = [block() for _ in range(depth)]
+
+    def __call__(self, latents, c, **kwargs):
+        for i, block in enumerate(self.blocks):
+            kwargs["layer_idx"] = i
+            latents = block(latents=latents, c=c, **kwargs)
+        return latents

models/simplefold/mlx/confidence_module.py

@@ -0,0 +1,84 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+#
+
+import mlx.nn as nn
+import mlx.core as mx
+
+
+def compute_aggregated_metric(logits, end=1.0):
+    """Compute the metric from the logits.
+
+    Parameters
+    ----------
+    logits : torch.Tensor
+        The logits of the metric
+    end : float
+        Max value of the metric, by default 1.0
+
+    Returns
+    -------
+    Tensor
+        The metric value
+
+    """
+    num_bins = logits.shape[-1]
+    bin_width = end / num_bins
+    bounds = mx.arange(start=0.5 * bin_width, stop=end, step=bin_width)
+    probs = mx.softmax(logits, axis=-1)
+    plddt = mx.sum(
+        probs * bounds.reshape(*((1,) * len(probs.shape[:-1])), *bounds.shape),
+        axis=-1,
+    )
+    return plddt
+
+
+class ConfidenceModule(nn.Module):
+    def __init__(
+        self,
+        hidden_size,
+        transformer_blocks=None,
+        num_plddt_bins=50,
+    ):
+        super().__init__()
+        self.transformer_blocks = transformer_blocks
+        self.to_plddt_logits = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size),
+            nn.LayerNorm(hidden_size),
+            nn.SiLU(),
+            nn.Linear(hidden_size, num_plddt_bins),
+        )
+
+    def __call__(
+        self,
+        latent,
+        feats,
+    ):
+        if self.transformer_blocks is not None:
+            token_pe_pos = mx.concatenate(
+                [
+                    feats["residue_index"][..., None].astype(mx.float32),  # (B, M, 1)
+                    feats["entity_id"][..., None].astype(mx.float32),  # (B, M, 1)
+                    feats["asym_id"][..., None].astype(mx.float32),  # (B, M, 1)
+                    feats["sym_id"][..., None].astype(mx.float32),  # (B, M, 1)
+                ],
+                axis=-1,
+            )
+
+            latent = self.transformer_blocks(
+                latents=latent,
+                c=None,
+                pos=token_pe_pos,
+            )
+
+        # Compute the pLDDT
+        plddt_logits = self.to_plddt_logits(latent)
+
+        # Compute the aggregated pLDDT
+        plddt = compute_aggregated_metric(plddt_logits)
+
+        return dict(
+            plddt=plddt,
+            plddt_logits=plddt_logits,
+        )

models/simplefold/mlx/esm_modules.py

@@ -0,0 +1,184 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+#
+
+# Started from https://github.com/facebookresearch/esm/tree/main,
+# licensed under MIT License, Copyright (c) Meta Platforms, Inc. and affiliates.
+
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from mlx.nn import LayerNorm as ESM1bLayerNorm
+from mlx.nn import gelu
+from .simplefold.mlx.esm_multihead_attention import MultiheadAttention  # noqa
+
+
+def symmetrize(x):
+    "Make layer symmetric in final two dimensions, used for contact prediction."
+    return x + mx.swapaxes(x, axis1=-1, axis2=-2)
+
+
+def apc(x):
+    "Perform average product correct, used for contact prediction."
+    a1 = x.sum(-1, keepdims=True)
+    a2 = x.sum(-2, keepdims=True)
+    a12 = x.sum((-1, -2), keepdims=True)
+
+    avg = a1 * a2
+    avg = mx.divide(avg, a12)
+    normalized = x - avg
+    return normalized
+
+
+class ESM1LayerNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-12, affine=True):
+        """Construct a layernorm layer in the TF style (eps inside the sqrt)."""
+        super().__init__()
+        self.hidden_size = (
+            (hidden_size,) if isinstance(hidden_size, int) else tuple(hidden_size)
+        )
+        self.eps = eps
+        self.affine = bool(affine)
+        if self.affine:
+            self.weight = mx.array(mx.ones(hidden_size))
+            self.bias = mx.array(mx.zeros(hidden_size))
+        else:
+            self.weight, self.bias = None, None
+
+    def __call__(self, x):
+        dims = tuple(-(i + 1) for i in range(len(self.hidden_size)))
+        means = x.mean(dims, keepdims=True)
+        x_zeromean = x - means
+        variances = x_zeromean.pow(2).mean(dims, keepdims=True)
+        x = x_zeromean / mx.sqrt(variances + self.eps)
+        if self.affine:
+            x = (self.weight * x) + self.bias
+        return x
+
+
+class TransformerLayer(nn.Module):
+
+    def __init__(
+        self,
+        embed_dim,
+        ffn_embed_dim,
+        attention_heads,
+        add_bias_kv=True,
+        use_esm1b_layer_norm=False,  # This is true in the implementation
+        use_rotary_embeddings: bool = False,  # This is true in the implementation
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.ffn_embed_dim = ffn_embed_dim
+        self.attention_heads = attention_heads
+        self.use_rotary_embeddings = use_rotary_embeddings
+        self._init_submodules(add_bias_kv, use_esm1b_layer_norm)
+
+    def _init_submodules(self, add_bias_kv, use_esm1b_layer_norm):
+        BertLayerNorm = ESM1bLayerNorm if use_esm1b_layer_norm else ESM1LayerNorm
+
+        self.self_attn = MultiheadAttention(
+            self.embed_dim,
+            self.attention_heads,
+            add_bias_kv=add_bias_kv,
+            add_zero_attn=False,
+            use_rotary_embeddings=self.use_rotary_embeddings,
+        )
+        self.self_attn_layer_norm = BertLayerNorm(self.embed_dim)
+
+        self.fc1 = nn.Linear(self.embed_dim, self.ffn_embed_dim)
+        self.fc2 = nn.Linear(self.ffn_embed_dim, self.embed_dim)
+
+        self.final_layer_norm = BertLayerNorm(self.embed_dim)
+
+    def __call__(
+        self,
+        x,
+        self_attn_mask=None,
+        self_attn_padding_mask=None,
+        need_head_weights=False,
+    ):
+        residual = x
+        x = self.self_attn_layer_norm(x)
+        x, attn = self.self_attn(
+            query=x,
+            key=x,
+            value=x,
+            key_padding_mask=self_attn_padding_mask,
+            need_weights=True,
+            need_head_weights=need_head_weights,
+            attn_mask=self_attn_mask,
+        )
+        x = residual + x
+
+        residual = x
+        x = self.final_layer_norm(x)
+        x = gelu(self.fc1(x))
+        x = self.fc2(x)
+        x = residual + x
+
+        return x, attn
+
+
+class RobertaLMHead(nn.Module):
+    """Head for masked language modeling."""
+
+    def __init__(self, embed_dim, output_dim, weight):
+        super().__init__()
+        self.dense = nn.Linear(embed_dim, embed_dim)
+        self.layer_norm = ESM1bLayerNorm(embed_dim)
+        self.weight = weight
+        self.bias = mx.array(mx.zeros(output_dim))
+
+    def __call__(self, features):
+        x = self.dense(features)
+        x = gelu(x)
+        x = self.layer_norm(x)
+        # project back to size of vocabulary with bias
+        x = mx.matmul(x, mx.swapaxes(self.weight, axis1=0, axis2=1)) + self.bias
+        return x
+
+
+class ContactPredictionHead(nn.Module):
+    """Performs symmetrization, apc, and computes a logistic regression on the output features"""
+
+    def __init__(
+        self,
+        in_features: int,
+        prepend_bos: bool,
+        append_eos: bool,
+        bias=True,
+        eos_idx: Optional[int] = None,
+    ):
+        super().__init__()
+        self.in_features = in_features
+        self.prepend_bos = prepend_bos
+        self.append_eos = append_eos
+        if append_eos and eos_idx is None:
+            raise ValueError(
+                "Using an alphabet with eos token, but no eos token was passed in."
+            )
+        self.eos_idx = eos_idx
+        self.regression = nn.Linear(in_features, 1, bias)
+        self.activation = nn.Sigmoid()
+
+    def __call__(self, tokens, attentions):
+        # remove eos token attentions
+        if self.append_eos:
+            eos_mask = tokens.ne(self.eos_idx)
+            eos_mask = eos_mask[:, None, ...] * eos_mask[:, :, None, ...]
+            attentions = attentions * eos_mask[:, None, None, :, :]
+            attentions = attentions[..., :-1, :-1]
+        # remove cls token attentions
+        if self.prepend_bos:
+            attentions = attentions[..., 1:, 1:]
+        batch_size, layers, heads, seqlen, _ = attentions.shape
+        attentions = attentions.reshape(batch_size, layers * heads, seqlen, seqlen)
+
+        # features: B x C x T x T
+        attentions = apc(symmetrize(attentions))
+        attentions = attentions.transpose(0, 2, 3, 1)
+        return self.activation(self.regression(attentions).squeeze(3))

models/simplefold/mlx/esm_multihead_attention.py

@@ -0,0 +1,224 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+#
+
+# Started from https://github.com/facebookresearch/esm/tree/main,
+# licensed under MIT License, Copyright (c) Meta Platforms, Inc. and affiliates.
+
+import mlx.nn as nn
+import mlx.core as mx
+from .simplefold.mlx.esm_rotary_embedding import RotaryEmbedding
+
+
+def utils_softmax(x, dim: int, onnx_trace: bool = False):
+    return mx.softmax(x.astype(mx.float32), axis=dim)
+
+
+def masked_fill_mlx(x, mask, value):
+    return mx.where(mask, value, x)
+
+
+class MultiheadAttention(nn.Module):
+    """Multi-headed attention.
+
+    See "Attention Is All You Need" for more details.
+    """
+
+    def __init__(
+        self,
+        embed_dim,
+        num_heads,
+        kdim=None,
+        vdim=None,
+        dropout=0.0,
+        bias=True,
+        add_bias_kv: bool = False,
+        add_zero_attn: bool = False,
+        self_attention: bool = False,
+        encoder_decoder_attention: bool = False,
+        use_rotary_embeddings: bool = False,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.kdim = kdim if kdim is not None else embed_dim
+        self.vdim = vdim if vdim is not None else embed_dim
+        self.qkv_same_dim = self.kdim == embed_dim and self.vdim == embed_dim
+
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.head_dim = embed_dim // num_heads
+        assert (
+            self.head_dim * num_heads == self.embed_dim
+        ), "embed_dim must be divisible by num_heads"
+        self.scaling = self.head_dim**-0.5
+
+        self.self_attention = self_attention
+
+        self.k_proj = nn.Linear(self.kdim, embed_dim, bias=bias)
+        self.v_proj = nn.Linear(self.vdim, embed_dim, bias=bias)
+        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+        if add_bias_kv:
+            self.bias_k = mx.array(1, 1, embed_dim)
+            self.bias_v = mx.array(1, 1, embed_dim)
+        else:
+            self.bias_k = self.bias_v = None
+
+        self.add_zero_attn = add_zero_attn
+
+        self.rot_emb = RotaryEmbedding(dim=self.head_dim)
+
+        self.enable_torch_version = False
+
+    def __call__(
+        self,
+        query,
+        key,
+        value,
+        key_padding_mask=None,
+        incremental_state=None,
+        need_weights=True,
+        static_kv=False,
+        attn_mask=None,
+        before_softmax=False,
+        need_head_weights=False,
+    ):
+        """Input shape: Time x Batch x Channel
+
+        Args:
+            key_padding_mask (ByteTensor, optional): mask to exclude
+                keys that are pads, of shape `(batch, src_len)`, where
+                padding elements are indicated by 1s.
+            need_weights (bool, optional): return the attention weights,
+                averaged over heads (default: False).
+            attn_mask (ByteTensor, optional): typically used to
+                implement causal attention, where the mask prevents the
+                attention from looking forward in time (default: None).
+            before_softmax (bool, optional): return the raw attention
+                weights and values before the attention softmax.
+            need_head_weights (bool, optional): return the attention
+                weights for each head. Implies *need_weights*. Default:
+                return the average attention weights over all heads.
+        """
+
+        tgt_len, bsz, embed_dim = query.shape
+        assert embed_dim == self.embed_dim
+        assert list(query.shape) == [tgt_len, bsz, embed_dim]
+
+        if self.self_attention:
+            q = self.q_proj(query)
+            k = self.k_proj(query)
+            v = self.v_proj(query)
+        else:
+            assert key is not None and value is not None
+            q = self.q_proj(query)
+            k = self.k_proj(key)
+            v = self.v_proj(value)
+        q *= self.scaling
+
+        if self.bias_k is not None:
+            assert self.bias_v is not None
+
+            # TODO: mlx not support array repeat or new_zeros
+            k = mx.concatenate([k, mx.tile(self.bias_k, (1, bsz, 1))])
+            v = mx.concatenate([v, mx.tile(self.bias_v, (1, bsz, 1))])
+            if attn_mask is not None:
+                attn_mask = mx.concatenate(
+                    [
+                        attn_mask,
+                        mx.zeros((attn_mask.shape[0], 1), dtype=attn_mask.dtype),
+                    ],
+                    axis=1,
+                )
+            if key_padding_mask is not None:
+                key_padding_mask = mx.concatenate(
+                    [
+                        key_padding_mask,
+                        mx.zeros(
+                            (key_padding_mask.shape[0], 1), dtype=key_padding_mask.dtype
+                        ),
+                    ],
+                    axis=1,
+                )
+
+        q = mx.swapaxes(
+            mx.contiguous(q).reshape(tgt_len, bsz * self.num_heads, self.head_dim),
+            axis1=0,
+            axis2=1,
+        )
+
+        if k is not None:
+            k = mx.swapaxes(
+                mx.contiguous(k).reshape(-1, bsz * self.num_heads, self.head_dim),
+                axis1=0,
+                axis2=1,
+            )
+
+        if v is not None:
+            v = mx.swapaxes(
+                mx.contiguous(v).reshape(-1, bsz * self.num_heads, self.head_dim),
+                axis1=0,
+                axis2=1,
+            )
+
+        assert k is not None
+        src_len = k.shape[1]
+
+        # This is part of a workaround to get around fork/join parallelism
+        # not supporting Optional types.
+
+        if key_padding_mask is not None and key_padding_mask.ndim == 0:
+            key_padding_mask = None
+
+        if key_padding_mask is not None:
+            assert key_padding_mask.shape[0] == bsz
+            assert key_padding_mask.shape[1] == src_len
+
+        if self.rot_emb:
+
+            q, k = self.rot_emb(q, k)
+
+        attn_weights = mx.matmul(q, mx.swapaxes(k, axis1=1, axis2=2))
+        attn_weights = MultiheadAttention.apply_sparse_mask(
+            attn_weights, tgt_len, src_len, bsz
+        )
+
+        assert list(attn_weights.shape) == [bsz * self.num_heads, tgt_len, src_len]
+
+        if attn_mask is not None:
+            attn_mask = attn_mask[None, ...]
+            attn_weights += attn_mask
+
+        if key_padding_mask is not None:
+            # don't attend to padding symbols
+            attn_weights = attn_weights.reshape(bsz, self.num_heads, tgt_len, src_len)
+            attn_weights = masked_fill_mlx(
+                attn_weights,
+                (key_padding_mask[:, None, None, ...] == 1.0),
+                float("-inf"),
+            )
+            attn_weights = attn_weights.reshape(bsz * self.num_heads, tgt_len, src_len)
+
+        attn_weights_float = utils_softmax(attn_weights, dim=-1, onnx_trace=False)
+        attn_weights = attn_weights_float.astype(attn_weights.dtype)
+
+        attn_probs = attn_weights.astype(attn_weights.dtype)
+        assert v is not None
+        attn = mx.matmul(attn_probs, v)
+
+        assert list(attn.shape) == [bsz * self.num_heads, tgt_len, self.head_dim]
+
+        attn = mx.contiguous(mx.swapaxes(attn, axis1=0, axis2=1)).reshape(
+            tgt_len, bsz, embed_dim
+        )
+        attn = self.out_proj(attn)
+
+        attn_weights = None
+
+        return attn, attn_weights
+
+    def apply_sparse_mask(attn_weights, tgt_len: int, src_len: int, bsz: int):
+        return attn_weights

models/simplefold/mlx/esm_network.py

@@ -0,0 +1,160 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+#
+
+# Started from https://github.com/facebookresearch/esm/tree/main,
+# licensed under MIT License, Copyright (c) Meta Platforms, Inc. and affiliates.
+
+from typing import Union
+import mlx.core as mx
+import mlx.nn as nn
+
+from .simplefold.mlx.esm_modules import (
+    ContactPredictionHead,
+    ESM1bLayerNorm,
+    RobertaLMHead,
+    TransformerLayer,
+)
+from onescience.datapipes.esm import Alphabet
+
+
+def masked_fill_mlx(x, mask, value):
+    return mx.where(mask, value, x)
+
+
+class ESM2(nn.Module):
+    def __init__(
+        self,
+        num_layers: int = 33,
+        embed_dim: int = 1280,
+        attention_heads: int = 20,
+        alphabet: Union[Alphabet, str] = "ESM-1b",
+        token_dropout: bool = True,
+    ):
+        super().__init__()
+        self.num_layers = num_layers
+        self.embed_dim = embed_dim
+        self.attention_heads = attention_heads
+        if not isinstance(alphabet, Alphabet):
+            alphabet = Alphabet.from_architecture(alphabet)
+        self.alphabet = alphabet
+        self.alphabet_size = len(alphabet)
+        self.padding_idx = alphabet.padding_idx
+        self.mask_idx = alphabet.mask_idx
+        self.cls_idx = alphabet.cls_idx
+        self.eos_idx = alphabet.eos_idx
+        self.prepend_bos = alphabet.prepend_bos
+        self.append_eos = alphabet.append_eos
+        self.token_dropout = token_dropout
+
+        self._init_submodules()
+
+    def _init_submodules(self):
+        self.embed_scale = 1
+        self.embed_tokens = mx.zeros((self.alphabet_size, self.embed_dim))
+
+        self.layers = [
+            TransformerLayer(
+                self.embed_dim,
+                4 * self.embed_dim,
+                self.attention_heads,
+                add_bias_kv=False,
+                use_esm1b_layer_norm=True,
+                use_rotary_embeddings=True,
+            )
+            for _ in range(self.num_layers)
+        ]
+
+        self.contact_head = ContactPredictionHead(
+            self.num_layers * self.attention_heads,
+            self.prepend_bos,
+            self.append_eos,
+            eos_idx=self.eos_idx,
+        )
+        self.emb_layer_norm_after = ESM1bLayerNorm(self.embed_dim)
+
+        self.lm_head = RobertaLMHead(
+            embed_dim=self.embed_dim,
+            output_dim=self.alphabet_size,
+            weight=self.embed_tokens,
+        )
+
+    def __call__(
+        self, tokens, repr_layers=[], need_head_weights=False, return_contacts=False
+    ):
+        if return_contacts:
+            need_head_weights = True
+
+        assert tokens.ndim == 2
+        padding_mask = mx.equal(tokens, self.padding_idx)  # B, T
+
+        x = self.embed_scale * self.embed_tokens[tokens, :]
+
+        if self.token_dropout:
+            x = masked_fill_mlx(x, (tokens == self.mask_idx)[..., None], 0.0)
+            # x: B x T x C
+            mask_ratio_train = 0.15 * 0.8
+            src_lengths = (~padding_mask).sum(axis=-1)
+            mask_ratio_observed = (tokens == self.mask_idx).sum(axis=-1).astype(
+                x.dtype
+            ) / src_lengths
+            x = x * (1 - mask_ratio_train) / (1 - mask_ratio_observed)[:, None, None]
+
+        if padding_mask is not None:
+            x = x * (1 - padding_mask[..., None].astype(x.dtype))
+
+        repr_layers = set(repr_layers)
+        hidden_representations = {}
+        if 0 in repr_layers:
+            hidden_representations[0] = x
+
+        if need_head_weights:
+            attn_weights = []
+
+        # (B, T, E) => (T, B, E)
+        x = mx.swapaxes(x, axis1=0, axis2=1)
+
+        if not padding_mask.any():
+            padding_mask = None
+
+        for layer_idx, layer in enumerate(self.layers):
+            x, attn = layer(
+                x,
+                self_attn_padding_mask=padding_mask,
+                need_head_weights=need_head_weights,
+            )
+
+            if (layer_idx + 1) in repr_layers:
+                hidden_representations[layer_idx + 1] = mx.swapaxes(x, axis1=0, axis2=1)
+            if need_head_weights:
+                # (H, B, T, T) => (B, H, T, T)
+                attn_weights.append(mx.swapaxes(attn, axis1=1, axis2=0))
+
+        x = self.emb_layer_norm_after(x)
+        x = mx.swapaxes(x, axis1=0, axis2=1)  # (T, B, E) => (B, T, E)
+
+        # last hidden representation should have layer norm applied
+        if (layer_idx + 1) in repr_layers:
+            hidden_representations[layer_idx + 1] = x
+        x = self.lm_head(x)
+
+        result = {"logits": x, "representations": hidden_representations}
+        if need_head_weights:
+            # attentions: B x L x H x T x T
+            attentions = mx.stack(attn_weights, axis=1)
+            if padding_mask is not None:
+                attention_mask = 1 - padding_mask.astype(attentions.dtype)
+                attention_mask = (
+                    attention_mask[:, None, ...] * attention_mask[:, :, None, ...]
+                )
+                attentions = attentions * attention_mask[:, None, None, :, :]
+            result["attentions"] = attentions
+            if return_contacts:
+                contacts = self.contact_head(tokens, attentions)
+                result["contacts"] = contacts
+
+        return result
+
+    def predict_contacts(self, tokens):
+        return self(tokens, return_contacts=True)["contacts"]

models/simplefold/mlx/esm_rotary_embedding.py

@@ -0,0 +1,74 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+#
+
+# Started from https://github.com/facebookresearch/esm/tree/main,
+# licensed under MIT License, Copyright (c) Meta Platforms, Inc. and affiliates.
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+def rotate_half(x):
+    x1, x2 = x.split(2, axis=-1)
+    return mx.concatenate([-x2, x1], axis=-1)
+
+
+def apply_rotary_pos_emb(x, cos, sin):
+    cos = cos[:, : x.shape[-2], :]
+    sin = sin[:, : x.shape[-2], :]
+
+    return (x * cos) + (rotate_half(x) * sin)
+
+
+class RotaryEmbedding(nn.Module):
+    """
+    The rotary position embeddings from RoFormer_ (Su et. al).
+    A crucial insight from the method is that the query and keys are
+    transformed by rotation matrices which depend on the relative positions.
+    Other implementations are available in the Rotary Transformer repo_ and in
+    GPT-NeoX_, GPT-NeoX was an inspiration
+    .. _RoFormer: https://arxiv.org/abs/2104.09864
+    .. _repo: https://github.com/ZhuiyiTechnology/roformer
+    .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox
+    .. warning: Please note that this embedding is not registered on purpose, as it is transformative
+        (it does not create the embedding dimension) and will likely be picked up (imported) on a ad-hoc basis
+    """
+
+    def __init__(self, dim: int, *_, **__):
+        super().__init__()
+        # Generate and save the inverse frequency buffer (non trainable)
+        self.inv_freq = mx.array(
+            1.0 / (10000 ** (mx.arange(0, dim, 2).astype(mx.float32) / dim))
+        )
+
+        self._seq_len_cached = None
+        self._cos_cached = None
+        self._sin_cached = None
+
+    def _update_cos_sin_tables(self, x, seq_dimension=1):
+        seq_len = x.shape[seq_dimension]
+
+        # Reset the tables if the sequence length has changed,
+        # or if we're on a new device (possibly due to tracing for instance)
+        if seq_len != self._seq_len_cached:
+            self._seq_len_cached = seq_len
+            t = mx.arange(x.shape[seq_dimension]).astype(self.inv_freq.dtype)
+            freqs = mx.einsum("i,j->ij", t, self.inv_freq)
+            emb = mx.concatenate([freqs, freqs], axis=-1)
+
+            self._cos_cached = emb.cos()[None, :, :]
+            self._sin_cached = emb.sin()[None, :, :]
+
+        return self._cos_cached, self._sin_cached
+
+    def __call__(self, q, k):
+        self._cos_cached, self._sin_cached = self._update_cos_sin_tables(
+            k, seq_dimension=-2
+        )
+
+        return (
+            apply_rotary_pos_emb(q, self._cos_cached, self._sin_cached),
+            apply_rotary_pos_emb(k, self._cos_cached, self._sin_cached),
+        )

models/simplefold/mlx/layers.py

@@ -0,0 +1,263 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+#
+
+import math
+from einops.array_api import rearrange
+from operator import __add__
+import mlx.core as mx
+import mlx.nn as nn
+
+
+def modulate(x, shift, scale):
+    return x * (1 + mx.expand_dims(scale, axis=1)) + mx.expand_dims(shift, axis=1)
+
+
+#################################################################################
+#                            Attention Layers                                  #
+#################################################################################
+
+
+class SelfAttentionLayer(nn.Module):
+    def __init__(
+        self,
+        hidden_size,
+        num_heads=8,
+        qkv_bias=False,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+        use_bias=True,
+        qk_norm=True,
+        pos_embedder=None,
+        linear_target: nn.Module = nn.Linear,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = hidden_size // num_heads
+        # NOTE scale factor was wrong in my original version,
+        # can set manually to be compat with prev weights
+        self.scale = qk_scale or head_dim**-0.5
+
+        self.qkv = linear_target(hidden_size, hidden_size * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = linear_target(hidden_size, hidden_size, bias=use_bias)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        self.q_norm = nn.RMSNorm(head_dim, eps=1e-8) if qk_norm else nn.Identity()
+        self.k_norm = nn.RMSNorm(head_dim, eps=1e-8) if qk_norm else nn.Identity()
+
+        self.pos_embedder = pos_embedder
+
+    def __call__(self, x, **kwargs):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
+        pos = kwargs.get("pos")
+
+        qkv = rearrange(qkv, "b n t h c -> t b h n c")
+        q, k, v = (
+            qkv[0],
+            qkv[1],
+            qkv[2],
+        )
+
+        q, k = self.q_norm(q), self.k_norm(k)
+
+        if self.pos_embedder and pos is not None:
+            q, k = self.pos_embedder(q, k, pos)
+
+        attn = (q @ k.swapaxes(axis1=-2, axis2=-1)) * self.scale
+        attn = mx.softmax(attn, axis=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).swapaxes(axis1=1, axis2=2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class EfficientSelfAttentionLayer(SelfAttentionLayer):
+    """Adapted from https://github.com/facebookresearch/dinov2/blob/main/dinov2/layers/attention.py"""
+
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+
+    def __call__(self, x, **kwargs):
+        B, N, C = x.shape
+        attn_mask = kwargs.get("attention_mask")
+        pos = kwargs.get("pos")
+
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
+        qkv = rearrange(qkv, "b n t h c -> t b h n c")
+
+        q, k, v = (
+            qkv[0],
+            qkv[1],
+            qkv[2],
+        )
+        if attn_mask is not None:
+            attn_mask = attn_mask.astype(q.dtype)
+
+        # if self.pos_embedder and pos is not None:
+        q, k = self.pos_embedder(q, k, pos)
+
+        q, k = self.q_norm(q), self.k_norm(k)
+
+        x = mx.fast.scaled_dot_product_attention(
+            q, k, v, mask=attn_mask, scale=1.0 / mx.sqrt(q.shape[-1])
+        )
+
+        x = x.swapaxes(axis1=1, axis2=2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+
+        return_attn = kwargs.get("return_attn", False)
+        if return_attn:
+            attn = (q @ k.swapaxes(axis1=-2, axis2=-1)) * self.scale
+            attn = attn.softmax(axis=-1)
+            return x, attn
+
+        return x, None
+
+
+def exists(val) -> bool:
+    """returns whether val is not none"""
+    return val is not None
+
+
+def default(x, y):
+    """returns x if it exists, otherwise y"""
+    return x if exists(x) else y
+
+
+#################################################################################
+#                              FeedForward Layer                                #
+#################################################################################
+
+
+class SwiGLUFeedForward(nn.Module):
+    def __init__(self, dim, hidden_dim, multiple_of=256):
+        super().__init__()
+        hidden_dim = int(2 * hidden_dim / 3)
+        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
+
+        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, dim, bias=True)
+        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
+
+    def __call__(self, x):
+        return self.w2(nn.silu(self.w1(x)) * self.w3(x))
+
+
+#################################################################################
+#                               Utility Layers                                  #
+#################################################################################
+
+
+class TimestepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations.
+    """
+
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        nn.init.normal(self.mlp.layers[0].weight, std=0.02)
+        nn.init.normal(self.mlp.layers[2].weight, std=0.02)
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        """
+        Create sinusoidal timestep embeddings.
+        :param t: a 1-D Tensor of N indices, one per batch element.
+                          These may be fractional.
+        :param dim: the dimension of the output.
+        :param max_period: controls the minimum frequency of the embeddings.
+        :return: an (N, D) Tensor of positional embeddings.
+        """
+        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+        half = dim // 2
+        freqs = mx.exp(
+            -math.log(max_period)
+            * mx.arange(start=0, stop=half, dtype=mx.float32)
+            / half
+        )
+        args = t[:, None].astype(mx.float32) * freqs[None]
+        embedding = mx.concatenate([mx.cos(args), mx.sin(args)], axis=-1)
+        if dim % 2:
+            embedding = mx.concatenate(
+                [embedding, mx.zeros_like(embedding[:, :1])], axis=-1
+            )
+        return embedding
+
+    def __call__(self, t):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+class ConditionEmbedder(nn.Module):
+    """
+    Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance.
+    """
+    def __init__(self, input_dim, hidden_size, dropout_prob):
+        super().__init__()
+        self.proj = nn.Sequential(
+                nn.Linear(input_dim, hidden_size),
+                nn.LayerNorm(hidden_size),
+                nn.SiLU(),
+            )
+        self.dropout_prob = dropout_prob
+        self.null_token = mx.zeros(input_dim)
+
+    def token_drop(self, cond, force_drop_ids=None):
+        """
+        cond: (B, N, D)
+        Drops conditions to enable classifier-free guidance.
+        """
+        if force_drop_ids is None:
+            drop_ids = mx.random.uniform(cond.shape[0]) < self.dropout_prob
+        else:
+            drop_ids = force_drop_ids
+        cond[drop_ids] = self.null_token[None, None, :]
+        return cond
+
+    def __call__(self, cond, train, force_drop_ids=None):
+        use_dropout = self.dropout_prob > 0
+        if (train and use_dropout) or (force_drop_ids is not None):
+            cond = self.token_drop(cond, force_drop_ids)
+        embeddings = self.proj(cond)
+        return embeddings
+
+
+class FinalLayer(nn.Module):
+    """
+    The final layer of DiT.
+    """
+
+    def __init__(self, hidden_size, out_channels, c_dim=None):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, affine=False, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, out_channels, bias=True)
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(), nn.Linear(c_dim, 2 * hidden_size, bias=True)
+        )
+
+    def __call__(self, x, c):
+        shift, scale = self.adaLN_modulation(c).split(2, axis=1)
+        x = modulate(self.norm_final(x), shift, scale)
+        x = self.linear(x)
+        return x

models/simplefold/mlx/pos_embed.py

@@ -0,0 +1,204 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+#
+
+import mlx.core as mx
+import mlx.nn as nn
+import math
+from einops.array_api import rearrange
+import torch
+import numpy as np
+
+
+class AbsolutePositionEncoding(nn.Module):
+    def __init__(self, in_dim, embed_dim, include_input=False):
+        super().__init__()
+        self.in_dim = in_dim
+        self.hidden_dim = embed_dim
+        self.include_input = include_input
+        assert embed_dim % in_dim == 0, "embed_dim must be divisible by in_dim"
+        self.embed_dim = embed_dim + in_dim if include_input else embed_dim
+
+    def __call__(self, pos):
+        pos_embs = []
+        for i in range(self.in_dim):
+            pe = self.get_1d_pos_embed(pos[..., i])
+            pos_embs.append(pe)
+        if self.include_input:
+            pos_embs.append(pos)
+        pos_embs = mx.concatenate(pos_embs, axis=-1)
+        return pos_embs
+
+    def get_1d_pos_embed(self, pos):
+        """
+        https://github.com/facebookresearch/DiT/blob/main/models.py#L303
+        """
+        embed_dim = self.hidden_dim // (self.in_dim * 2)
+        omega = 2 ** mx.linspace(0, math.log(224, 2) - 1, embed_dim).astype(mx.float32)
+        omega *= math.pi
+
+        if len(pos.shape) == 1:
+            out = mx.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
+        elif len(pos.shape) == 2:
+            out = mx.einsum("nm,d->nmd", pos, omega)
+
+        emb_sin = mx.sin(out)
+        emb_cos = mx.cos(out)  # (*, M, D/2)
+        emb = mx.concatenate([emb_sin, emb_cos], axis=-1)  # (*, M, D)
+        return emb
+
+
+class FourierPositionEncoding(nn.Module):
+    def __init__(
+        self,
+        in_dim: int,
+        include_input: bool = False,
+        min_freq_log2: float = 0,
+        max_freq_log2: float = 12,
+        num_freqs: int = 32,
+        log_sampling: bool = True,
+    ):
+        super().__init__()
+        self.in_dim = in_dim
+        self.include_input = include_input
+        self.min_freq_log2 = min_freq_log2
+        self.max_freq_log2 = max_freq_log2
+        self.num_freqs = num_freqs
+        self.log_sampling = log_sampling
+        self.create_embedding_fn()
+
+    def create_embedding_fn(self):
+        d = self.in_dim
+        dim_out = 0
+        if self.include_input:
+            dim_out += d
+
+        min_freq = self.min_freq_log2
+        max_freq = self.max_freq_log2
+        N_freqs = self.num_freqs
+
+        if self.log_sampling:
+            freq_bands = 2.0 ** mx.linspace(min_freq, max_freq, num=N_freqs)  # (nf,)
+        else:
+            freq_bands = mx.linspace(2.0**min_freq, 2.0**max_freq, num=N_freqs)  # (nf,)
+
+        assert mx.isfinite(
+            freq_bands
+        ).all(), f"nan: {mx.isnan(freq_bands).any()} inf: {mx.isinf(freq_bands).any()}"
+
+        self.freq_bands = freq_bands
+        self.embed_dim = dim_out + d * self.freq_bands.size * 2
+
+    def __call__(
+        self,
+        pos,
+    ):
+        """
+        Get the positional encoding for each coordinate.
+        Args:
+            pos:
+                (*, in_dim)
+        Returns:
+            out:
+                (*, in_dimitional_encoding)
+        """
+
+        out = []
+        if self.include_input:
+            out = [pos]  # (*, in_dim)
+
+        pos = pos[..., None] * self.freq_bands  # (*b, d, nf)
+
+        out += [
+            mx.sin(pos).flatten(start_axis=-2),  # (*b, d*nf)
+            mx.cos(pos).flatten(start_axis=-2),  # (*b, d*nf)
+        ]
+
+        out = mx.concatenate(out, axis=-1)  # (*b, 2 * in_dim * nf (+ in_dim))
+        return out
+
+
+def compute_axial_cis(
+    ts,
+    in_dim: int,
+    dim: int,
+    theta: float = 100.0,
+):
+    B, N, D = ts.shape
+    freqs_all = []
+    interval = 2 * in_dim
+    for i in range(in_dim):
+        freq = 1.0 / (
+            theta
+            ** (
+                mx.arange(0, dim, interval)[: (dim // interval)].astype(mx.float32)
+                / dim
+            )
+        )
+        t = ts[..., i].flatten()
+        freq_i = mx.outer(t, freq)
+        freq_cis_i = polar(mx.ones_like(freq_i), freq_i)
+        freq_cis_i = freq_cis_i.reshape((B, N, -1))
+        freqs_all.append(freq_cis_i)
+    freqs_cis = mx.concatenate(freqs_all, axis=-1)
+    return freqs_cis
+
+
+def polar(a, b):
+    return a * mx.exp(1j * b)
+
+
+def view_as_complex(x):
+    x = x.astype(mx.float32)
+    x = x.reshape(*x.shape[:-1], -1, 2)  # (..., dim/2, 2)
+    return x[..., 0] + 1j * x[..., 1]  # real, imag
+
+
+def view_as_real(input):
+    return mx.stack([input.real, input.imag], axis=-1)  # (..., dim)
+
+
+def apply_rotary_emb(xq: mx.array, xk: mx.array, freqs_cis: mx.array):
+    # xq, xk: shape (..., dim)
+    # freqs_cis: shape (..., dim // 2, 2) where last dim = [cos, sin]
+    xq_ = view_as_complex(xq)
+    xk_ = view_as_complex(xk)
+
+    # Apply complex multiplication
+    xq_out = xq_ * freqs_cis
+    xk_out = xk_ * freqs_cis
+
+    # Reconstruct to original shape
+    xq_out = view_as_real(xq_out).flatten(start_axis=3)
+    xk_out = view_as_real(xk_out).flatten(start_axis=3)
+
+    return xq_out.astype(xq.dtype), xk_out.astype(xk.dtype)
+
+
+class AxialRotaryPositionEncoding(nn.Module):
+    def __init__(
+        self,
+        in_dim,
+        embed_dim,
+        num_heads,
+        base=100.0,
+    ):
+        super().__init__()
+        self.in_dim = in_dim
+        self.num_heads = num_heads
+        self.embed_dim = embed_dim // num_heads
+        self.base = base
+
+    def __call__(self, xq, xk, pos):
+        """
+        xq: [B, H, N, D]
+        xk: [B, H, N, D]
+        pos: [B, N, in_dim]
+        """
+        if pos.ndim == 2:
+            pos = pos[..., None]
+        freqs_cis = compute_axial_cis(pos, self.in_dim, self.embed_dim, self.base)
+        freqs_cis = mx.expand_dims(freqs_cis, axis=1)
+
+        return apply_rotary_emb(xq, xk, freqs_cis)

models/simplefold/mlx/sampler.py

@@ -0,0 +1,112 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+#
+
+import mlx.core as mx
+from tqdm import tqdm
+from einops.array_api import repeat
+from onescience.utils.simplefold.mlx_utils import center_random_augmentation
+
+
+def logspace(start, end, steps, base=10.0, dtype=mx.float32):
+    # create a linear space between start and end
+    lin = mx.linspace(start, end, steps, dtype=dtype)
+    # raise base to that power
+    return mx.power(mx.array(base, dtype=dtype), lin)
+
+
+class EMSampler:
+    """
+    A Euler-Maruyama solver for SDEs.
+    """
+
+    def __init__(
+        self,
+        num_timesteps=500,
+        t_start=1e-4,
+        tau=0.3,
+        log_timesteps=False,
+        w_cutoff=0.99,
+    ):
+        self.num_timesteps = num_timesteps
+        self.log_timesteps = log_timesteps
+        self.t_start = t_start
+        self.tau = tau
+        self.w_cutoff = w_cutoff
+
+        if self.log_timesteps:
+            t = 1.0 - logspace(-2, 0, steps=self.num_timesteps + 1)[::-1, ...]
+            t = t - mx.min(t)
+            t = t / mx.max(t)
+            self.steps = mx.clip(t, a_min=self.t_start, a_max=1.0)
+        else:
+            self.steps = mx.linspace(self.t_start, 1.0, num=self.num_timesteps + 1)
+
+    def diffusion_coefficient(self, t, eps=0.01):
+        # determine diffusion coefficient
+        w = (1.0 - t) / (t + eps)
+        if t >= self.w_cutoff:
+            w = 0.0
+        return w
+
+    def euler_maruyama_step(
+        self,
+        model_fn,
+        flow,
+        y,
+        t,
+        t_next,
+        batch,
+    ):
+        dt = t_next - t
+        eps = mx.random.normal(y.shape)
+
+        y = center_random_augmentation(
+            y,
+            batch["atom_pad_mask"],
+            augmentation=False,
+            centering=True,
+        )
+
+        batched_t = repeat(t, " -> b", b=y.shape[0])
+        velocity = model_fn(
+            noised_pos=y,
+            t=batched_t,
+            feats=batch,
+        )["predict_velocity"]
+
+        score = flow.compute_score_from_velocity(velocity, y, t)
+
+        diff_coeff = self.diffusion_coefficient(t)
+        drift = velocity + diff_coeff * score
+        mean_y = y + drift * dt
+        y_sample = mean_y + mx.sqrt(2.0 * dt * diff_coeff * self.tau) * eps
+
+        return y_sample
+
+    def sample(self, model_fn, flow, noise, batch):
+        sampling_timesteps = self.num_timesteps
+        steps = self.steps
+        y_sampled = noise
+        feats = batch
+
+        for i in tqdm(
+            range(sampling_timesteps),
+            desc="Sampling",
+            total=sampling_timesteps,
+        ):
+            t = steps[i]
+            t_next = steps[i + 1]
+
+            y_sampled = self.euler_maruyama_step(
+                model_fn,
+                flow,
+                y_sampled,
+                t,
+                t_next,
+                feats,
+            )
+            mx.eval(y_sampled)
+
+        return {"denoised_coords": y_sampled}

models/simplefold/simplefold.py

@@ -0,0 +1,781 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+#
+
+import os
+import shutil
+import copy
+import numpy as np
+from pathlib import Path
+from einops import repeat
+
+import torch
+import torch.nn.functional as F
+from torch.optim.optimizer import Optimizer
+import torch.distributed as dist
+from torch.distributed.fsdp import FullyShardedDataParallel
+from torch.optim.swa_utils import AveragedModel
+from torch.nn.utils import clip_grad_norm_
+
+import lightning
+import lightning.pytorch as pl
+from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
+from fairscale.nn.wrap import enable_wrap, wrap
+
+from onescience.utils.simplefold.esm_utils import _af2_to_esm, esm_registry
+from onescience.datapipes.boltz_data_pipeline.types import Record, Structure
+from onescience.utils.simplefold.boltz_utils import (
+    weighted_rigid_align, 
+    center_random_augmentation,
+    process_structure, 
+    save_structure
+)
+
+
+def logit_normal_sample(n=1, m=0.0, s=1.0):
+    # Logit-Normal Sampling from https://arxiv.org/pdf/2403.03206.pdf
+    u = torch.randn(n) * s + m
+    t = 1 / (1 + torch.exp(-u))
+    return t
+
+
+def lddt_dist(dmat_predicted, dmat_true, mask, cutoff=15.0, per_atom=False):
+    # NOTE: the mask is a pairwise mask which should have the identity elements already masked out
+    # Compute mask over distances
+    dists_to_score = (dmat_true < cutoff).float() * mask
+    dist_l1 = torch.abs(dmat_true - dmat_predicted)
+
+    score = 0.25 * (
+        (dist_l1 < 0.5).float()
+        + (dist_l1 < 1.0).float()
+        + (dist_l1 < 2.0).float()
+        + (dist_l1 < 4.0).float()
+    )
+
+    # Normalize over the appropriate axes.
+    if per_atom:
+        mask_no_match = torch.sum(dists_to_score, dim=-1) != 0
+        norm = 1.0 / (1e-10 + torch.sum(dists_to_score, dim=-1))
+        score = norm * (1e-10 + torch.sum(dists_to_score * score, dim=-1))
+        return score, mask_no_match.float()
+    else:
+        norm = 1.0 / (1e-10 + torch.sum(dists_to_score, dim=(-2, -1)))
+        score = norm * (1e-10 + torch.sum(dists_to_score * score, dim=(-2, -1)))
+        total = torch.sum(dists_to_score, dim=(-1, -2))
+        return score, total
+
+
+class SimpleFold(pl.LightningModule):
+    def __init__(
+        self,
+        architecture,
+        processor,
+        loss,
+        path,
+        sampler,
+        optimizer=None,
+        scheduler=None,
+        plddt_module=None,
+        ema_decay=0.999,
+        esm_model="esm2_3B",
+        aa_bolt_link=None,
+        use_rigid_align=True,
+        smooth_lddt_loss_weight=1.0,
+        lddt_cutoff=15.0,
+        clip_grad_norm_val=None,
+        lddt_weight_schedule=False,
+        plddt_training=False,
+        sample_dir='artifacts/',
+    ):
+        super().__init__()
+        self.save_hyperparameters(logger=False)
+
+        self.model = architecture
+        self.model_ema = AveragedModel(
+            self.model,
+            multi_avg_fn=torch.optim.swa_utils.get_ema_multi_avg_fn(
+                self.hparams.ema_decay
+            ),
+            use_buffers=True,
+        )
+        self.model_ema.eval()
+
+        self.loss = loss
+        self.path = path
+        self.sampler = sampler
+
+        self.use_rigid_align = use_rigid_align
+        self.lddt_cutoff = lddt_cutoff
+        self.smooth_lddt_loss_weight = smooth_lddt_loss_weight
+        self.use_smooth_lddt_loss = smooth_lddt_loss_weight > 0.0
+        self.lddt_weight_schedule = lddt_weight_schedule
+        self.plddt_training = plddt_training
+        self.sample_dir = sample_dir
+
+        self.aa_bolt_link = aa_bolt_link
+        self.nval_steps = 0
+
+        try:
+            self.t_eps = self.sampler.t_eps
+        except AttributeError:
+            self.t_eps = 0.0
+
+        self.use_esm = esm_model is not None
+        if self.use_esm:
+            self.esm_model, self.esm_dict = esm_registry[esm_model]()
+            self.esm_model.eval()
+            self.af2_to_esm = _af2_to_esm(self.esm_dict)
+            print(f"Using ESM model: {esm_model}")
+        else:
+            self.esm_model = None
+            self.esm_dict = None
+            self.af2_to_esm = None
+
+        self.plddt_module = plddt_module
+        if self.plddt_training:
+            assert self.plddt_module is not None, "PLDDT module must be provided for PLDDT training"
+            self.model.eval()
+
+    def register(self, name, tensor):
+        self.register_buffer(name, tensor.type(torch.float32))
+
+    def loss_masking(self, loss, atom_mask):
+        loss_mask = repeat(atom_mask, "b s -> b s d", d=loss.shape[-1])
+        loss *= loss_mask
+
+        denom = torch.sum(atom_mask, -1, keepdim=True)
+        denom = denom.unsqueeze(-1)
+        loss = torch.sum(loss, dim=1, keepdim=True) / denom
+        return loss
+
+    def smooth_lddt_loss(
+        self, 
+        pred_coords,
+        true_coords,
+        # is_nucleotide,
+        coords_mask,
+        t,
+    ):
+        """Compute weighted alignment.
+
+        Parameters
+        ----------
+        pred_coords: torch.Tensor
+            The predicted atom coordinates
+        true_coords: torch.Tensor
+            The ground truth atom coordinates
+        coords_mask: torch.Tensor
+            The atoms mask
+
+        """
+        B, N, _ = true_coords.shape
+        true_dists = torch.cdist(true_coords, true_coords)
+
+        mask = (true_dists < self.lddt_cutoff).float()
+        mask = mask * (1 - torch.eye(pred_coords.shape[1], device=pred_coords.device))
+        mask = mask * (coords_mask.unsqueeze(-1) * coords_mask.unsqueeze(-2))
+
+        # Compute distances between all pairs of atoms
+        pred_dists = torch.cdist(pred_coords, pred_coords)
+        dist_diff = torch.abs(true_dists - pred_dists)
+
+        # Compute epsilon values
+        eps = (
+            (
+                (
+                    F.sigmoid(0.5 - dist_diff)
+                    + F.sigmoid(1.0 - dist_diff)
+                    + F.sigmoid(2.0 - dist_diff)
+                    + F.sigmoid(4.0 - dist_diff)
+                )
+                / 4.0
+            )
+            .view(B, N, N)
+            .mean(dim=0)
+        )
+
+        # Calculate masked averaging
+        num = (eps * mask).sum(dim=(-1, -2))
+        den = mask.sum(dim=(-1, -2)).clamp(min=1)
+        lddt = num / den
+        if self.lddt_weight_schedule:
+            t_weight = 1 + 8 * torch.relu(t - 0.5)
+            lddt = (1.0 - lddt) * t_weight
+            return lddt.mean()
+        else:
+            return (1.0 - lddt.mean()) * self.smooth_lddt_loss_weight
+
+    def plddt_loss(
+        self,
+        pred_lddt,
+        pred_atom_coords,
+        true_atom_coords,
+        true_coords_resolved_mask,
+        feats,
+        # multiplicity=1,
+    ):
+        """Compute plddt loss.
+
+        Parameters
+        ----------
+        pred_lddt: torch.Tensor
+            The plddt logits
+        pred_atom_coords: torch.Tensor
+            The predicted atom coordinates
+        true_atom_coords: torch.Tensor
+            The atom coordinates after symmetry correction
+        true_coords_resolved_mask: torch.Tensor
+            The resolved mask after symmetry correction
+        feats: Dict[str, torch.Tensor]
+            Dictionary containing the model input
+
+        Returns
+        -------
+        torch.Tensor
+            Plddt loss
+
+        """
+
+        # extract necessary features
+        atom_mask = true_coords_resolved_mask
+
+        R_set_to_rep_atom = feats["r_set_to_rep_atom"].float()
+        # R_set_to_rep_atom = R_set_to_rep_atom.repeat_interleave(multiplicity, 0).float()
+
+        token_type = feats["mol_type"]
+        # token_type = token_type.repeat_interleave(multiplicity, 0)
+        # is_nucleotide_token = (token_type == const.chain_type_ids["DNA"]).float() + (
+        #     token_type == const.chain_type_ids["RNA"]
+        # ).float()
+
+        B = true_atom_coords.shape[0]
+
+        # atom_to_token = feats["atom_to_token"].float()
+        # atom_to_token = atom_to_token.repeat_interleave(multiplicity, 0)
+
+        token_to_rep_atom = feats["token_to_rep_atom"].float()
+        # token_to_rep_atom = token_to_rep_atom.repeat_interleave(multiplicity, 0)
+
+        true_token_coords = torch.bmm(token_to_rep_atom, true_atom_coords)
+        pred_token_coords = torch.bmm(token_to_rep_atom, pred_atom_coords)
+
+        # compute true lddt
+        true_d = torch.cdist(
+            true_token_coords,
+            torch.bmm(R_set_to_rep_atom, true_atom_coords),
+        )
+        pred_d = torch.cdist(
+            pred_token_coords,
+            torch.bmm(R_set_to_rep_atom, pred_atom_coords),
+        )
+
+        # compute mask
+        pair_mask = atom_mask.unsqueeze(-1) * atom_mask.unsqueeze(-2)
+        pair_mask = (
+            pair_mask
+            * (1 - torch.eye(pair_mask.shape[1], device=pair_mask.device))[None, :, :]
+        )
+        pair_mask = torch.einsum("bnm,bkm->bnk", pair_mask, R_set_to_rep_atom)
+        pair_mask = torch.bmm(token_to_rep_atom, pair_mask)
+        atom_mask = torch.bmm(token_to_rep_atom, atom_mask.unsqueeze(-1).float())
+        # is_nucleotide_R_element = torch.bmm(
+        #     R_set_to_rep_atom, torch.bmm(atom_to_token, is_nucleotide_token.unsqueeze(-1))
+        # ).squeeze(-1)
+        # cutoff = 15 + 15 * is_nucleotide_R_element.reshape(B, 1, -1).repeat(
+        #     1, true_d.shape[1], 1
+        # )
+
+        # compute lddt
+        target_lddt, mask_no_match = lddt_dist(
+            pred_d, true_d, pair_mask, cutoff=15.0, per_atom=True
+        )
+
+        # compute loss
+        num_bins = pred_lddt.shape[-1]
+        bin_index = torch.floor(target_lddt * num_bins).long()
+        bin_index = torch.clamp(bin_index, max=(num_bins - 1))
+        lddt_one_hot = F.one_hot(bin_index, num_classes=num_bins)
+        errors = -1 * torch.sum(
+            lddt_one_hot * F.log_softmax(pred_lddt, dim=-1),
+            dim=-1,
+        )
+        atom_mask = atom_mask.squeeze(-1)
+        loss = torch.sum(errors * atom_mask * mask_no_match, dim=-1) / (
+            1e-7 + torch.sum(atom_mask * mask_no_match, dim=-1)
+        )
+
+        # Average over the batch dimension
+        loss = torch.mean(loss)
+
+        self.log(
+            "loss/plddt",
+            loss.item(),
+            on_epoch=True,
+            logger=True,
+            prog_bar=True,
+            rank_zero_only=True,
+        )
+
+        return loss
+
+    def plddt_train_step(self, batch, batch_idx):
+        with torch.no_grad():
+            batch = self.processor.preprocess_training(
+                batch,
+                esm_model=self.esm_model,
+                esm_dict=self.esm_dict,
+                af2_to_esm=self.af2_to_esm,
+            )
+
+            noise = torch.randn_like(batch['coords']).to(self.device)
+
+            out_dict = self.sampler.sample(
+                self.model_ema.module.forward, self.path,
+                noise, batch
+            )
+            # out_dict = self.processor.postprocess(out_dict, batch)
+
+            # denoised_coords = center_of_mass_norm(
+            #     out_dict["denoised_coords"], batch['atom_pad_mask']
+            # )
+            # true_coords = center_of_mass_norm(
+            #     batch['coords'], batch['atom_pad_mask']
+            # )
+            denoised_coords = center_random_augmentation(
+                out_dict["denoised_coords"],
+                batch['atom_pad_mask'],
+                augmentation=False,
+                centering=True,
+            )
+            true_coords = center_random_augmentation(
+                batch['coords'],
+                batch['atom_pad_mask'],
+                augmentation=False,
+                centering=True,
+            )
+            out_dict["denoised_coords"] = denoised_coords * self.processor.scale
+            out_dict["coords"] = true_coords * self.processor.scale
+
+            out_dict["true_coords_resolved_mask"] = batch["atom_resolved_mask"]
+
+            t = torch.ones(batch['coords'].shape[0], device=self.device)
+            out_feat = self.model(denoised_coords, t, batch) # use unscaled coords
+
+        # Compute plddt loss
+        plddt_out_dict = self.plddt_module(
+            out_feat["latent"].detach(),
+            batch,
+        )
+
+        plddt_loss = self.plddt_loss(
+            plddt_out_dict["plddt_logits"],
+            out_dict["denoised_coords"],
+            out_dict["coords"],
+            out_dict["true_coords_resolved_mask"],
+            batch,
+        )
+
+        return plddt_loss
+
+    def flow_matching_train_step(self, batch, batch_idx):
+        batch = self.processor.preprocess_training(
+            batch,
+            esm_model=self.esm_model,
+            esm_dict=self.esm_dict,
+            af2_to_esm=self.af2_to_esm,
+        )
+
+        # timestep resampling
+        t_size = batch['coords'].shape[0]
+        t = 0.98 * logit_normal_sample(n=t_size, m=0.8, s=1.7) + 0.02 * torch.rand(t_size)
+        t = t.to(self.device)
+        t = t * (1 - 2 * self.t_eps) + self.t_eps
+
+        noise = torch.randn_like(batch['coords']).to(self.device)
+
+        _, y_t, v_t = self.path.interpolant(t, noise, batch["coords"])
+
+        out_dict = self.model(y_t, t, batch)
+
+        resolved_atom_mask = batch["atom_resolved_mask"].float()
+        align_weights = y_t.new_ones(y_t.shape[:2])
+
+        if self.use_rigid_align:
+            with torch.no_grad(), torch.autocast("cuda", enabled=False):
+                v_t = out_dict['predict_velocity'].detach().float()
+                denoised_coords = y_t + v_t * (1.0 - t[:, None, None])
+                coords = batch["coords"].detach().float()
+                coords_aligned = weighted_rigid_align(
+                    coords,
+                    denoised_coords.detach().float(),
+                    align_weights.detach().float(),
+                    mask=resolved_atom_mask.detach().float(),
+                )
+                _, _, v_t_aligned = self.path.interpolant(t, noise, coords_aligned)
+            target = v_t_aligned
+        else:
+            target = v_t
+
+        loss = F.mse_loss(out_dict['predict_velocity'], target, reduction='none')
+
+        loss_mask = resolved_atom_mask * align_weights
+        loss = self.loss_masking(loss, loss_mask)
+        loss = loss.mean()
+
+        self.log(
+            "loss/mse",
+            loss.item(),
+            on_epoch=True,
+            logger=True,
+            prog_bar=True,
+            rank_zero_only=True,
+        )
+
+        if self.use_smooth_lddt_loss:
+            # one-step Euler to get denoised coordinates
+            denoised_coords = y_t + \
+                out_dict['predict_velocity'] * (1.0 - t[:, None, None])
+
+            # rescale coordinates to angstroms
+            # denoised_coords = center_of_mass_norm(denoised_coords, batch['atom_pad_mask'])
+            # true_coords = center_of_mass_norm(batch['coords'], batch['atom_pad_mask'])
+            denoised_coords = center_random_augmentation(
+                denoised_coords,
+                batch['atom_pad_mask'],
+                augmentation=False,
+                centering=True,
+            )
+            true_coords = center_random_augmentation(
+                batch['coords'],
+                batch['atom_pad_mask'],
+                augmentation=False,
+                centering=True,
+            )
+            denoised_coords = denoised_coords * self.processor.scale
+            true_coords = true_coords * self.processor.scale
+
+            smooth_lddt_loss = self.smooth_lddt_loss(
+                denoised_coords,
+                true_coords,
+                resolved_atom_mask,
+                t,
+            )
+            loss += smooth_lddt_loss
+
+            self.log(
+                "loss/smooth_lddt",
+                smooth_lddt_loss.item(),
+                on_epoch=True,
+                logger=True,
+                prog_bar=True,
+                rank_zero_only=True,
+            )
+
+        self.log(
+            "loss/loss",
+            loss.item(),
+            on_epoch=True,
+            logger=True,
+            prog_bar=True,
+            rank_zero_only=True,
+        )
+
+        self.log(
+            "trainer/global_step",
+            self.global_step,
+            on_epoch=False,
+            logger=True,
+            prog_bar=False,
+            rank_zero_only=True,
+        )
+
+        self.global_training_step = self.trainer.global_step
+        self.epoch = self.trainer.current_epoch
+        self.world_size = self.trainer.world_size
+        return loss
+
+    def training_step(self, batch, batch_idx):
+        if self.plddt_training:
+            return self.plddt_train_step(batch, batch_idx)
+        else:
+            return self.flow_matching_train_step(batch, batch_idx)
+
+    @torch.no_grad()
+    def validation_step(self, batch, batch_idx):
+        # we skip validation step in training mode
+        return
+
+    def on_train_start(self):
+        global_seed = os.environ.get("PL_GLOBAL_SEED", 42)
+        pl.seed_everything(int(global_seed) + self.trainer.global_rank, True)
+        return
+
+    @torch.no_grad()
+    def predict_step(self, batch, batch_idx):
+        with torch.autocast(device_type='cuda', dtype=torch.float32):
+            batch = self.processor.preprocess_inference(
+                batch,
+                esm_model=self.esm_model,
+                esm_dict=self.esm_dict,
+                af2_to_esm=self.af2_to_esm,
+            )
+            num_repeats = batch.get("num_repeats", torch.tensor(1, device=self.device)).item()
+            multiplicity = batch["mol_type"].shape[0]
+            num_iter = np.ceil(num_repeats / multiplicity).astype(int)
+            print(f"Generating {num_repeats} samples with num_iter: {num_iter}, multiplicity: {multiplicity}")
+
+            # num_repeats is the total number of samples to generate for one protein
+            # multiplicity is the number of samples to generate at once
+
+            curr_idx = 0
+            for i in range(num_iter):
+                batch_in = copy.deepcopy(batch)
+                noise = torch.randn_like(batch_in['coords']).to(self.device)
+
+                out_dict = self.sampler.sample(
+                    self.model_ema.module.forward, self.path,
+                    noise, batch_in
+                )
+
+                if self.plddt_module is not None:
+                    denoised_coords = center_random_augmentation(
+                        out_dict["denoised_coords"],
+                        batch['atom_pad_mask'],
+                        augmentation=False,
+                        centering=True,
+                    )
+
+                    t = torch.ones(batch['coords'].shape[0], device=self.device)
+                    out_feat = self.model(denoised_coords, t, batch) # use unscaled coords
+
+                    plddt_out_dict = self.plddt_module(
+                        out_feat["latent"],
+                        batch_in,
+                    )
+                    plddts = plddt_out_dict["plddt"] * 100.0
+                else:
+                    plddts = None
+
+                out_dict = self.processor.postprocess(out_dict, batch_in)
+
+                record = Record(**batch_in['record'][0])
+                gt_coord = out_dict['coords']
+                sampled_coord = out_dict['denoised_coords']
+                pad_mask = batch_in['atom_pad_mask']
+
+                if num_repeats - curr_idx < multiplicity:
+                    curr_num_copies = num_repeats - curr_idx
+                    gt_coord = gt_coord[:curr_num_copies]
+                    sampled_coord = sampled_coord[:curr_num_copies]
+                    pad_mask = pad_mask[:curr_num_copies]
+                    plddts = plddts[:curr_num_copies] if plddts is not None else None
+                else:
+                    curr_num_copies = multiplicity
+
+                data_dir = self.trainer.datamodule.predict_dataloader().dataset.target_dir
+                sample_dir = self.sample_dir
+                path = Path(data_dir) / "structures" / f"{record.id}.npz"
+                structure: Structure = Structure.load(path)
+
+                for j in range(curr_num_copies):
+                    file_id = j + curr_idx
+                    try:
+                        sampled_structure = copy.deepcopy(structure)
+                        sampled_structure = process_structure(
+                            sampled_structure, sampled_coord[j], pad_mask[j], record
+                        )
+                        # save mmcif structure
+                        sampled_struct_dir = Path(sample_dir)
+                        outname = f"{record.id}_sampled_{str(file_id)}"
+                        save_structure(
+                            sampled_structure, sampled_struct_dir, outname, 
+                            plddts=plddts[j] if plddts is not None else None,
+                            output_format="mmcif",
+                        )
+                        # save pdb structure
+                        save_structure(
+                            sampled_structure, sampled_struct_dir, outname, 
+                            plddts=plddts[j] if plddts is not None else None,
+                            output_format="pdb",
+                        )
+                    except:
+                        print(f"Error processing {record.id}")
+                        continue
+
+                curr_idx += curr_num_copies
+
+            self.world_size = self.trainer.world_size
+
+        return
+
+    def on_predict_epoch_end(self):
+        dist.barrier() # wait for all processes to finish
+        return
+
+    def reset_esm(self, esm_model: str):
+        self.esm_model, self.esm_dict = esm_registry[esm_model]()
+        self.esm_model.eval()
+        self.af2_to_esm = _af2_to_esm(self.esm_dict)
+        self.esm_model = self.esm_model.to(self.device)
+        self.af2_to_esm = self.af2_to_esm.to(self.device)
+        print(f"Successfully reset ESM model {esm_model}")
+
+    def setup(self, stage: str) -> None:
+        """Lightning hook that is called at the beginning of fit (train + validate), validate,
+        test, or predict.
+
+        This is a good hook when you need to build models dynamically or adjust something about
+        them. This hook is called on every process when using DDP.
+
+        :param stage: Either `"fit"`, `"validate"`, `"test"`, or `"predict"`.
+        """
+        self.processor = self.hparams.processor(device=self.device)
+
+        if self.use_esm:
+            if stage == "fit" and not isinstance(
+                self.trainer.strategy, lightning.pytorch.strategies.fsdp.FSDPStrategy
+            ):
+                # initialize the model with FSDP wrapper
+                fsdp_params = dict(
+                    mixed_precision=True,
+                    flatten_parameters=True,
+                    state_dict_device=torch.device("cpu"),  # reduce GPU mem usage
+                    cpu_offload=True,  # enable cpu offloading
+                    fp32_reduce_scatter=True,  # use fp32 reduce scatter
+                )
+                with enable_wrap(wrapper_cls=FSDP, **fsdp_params):
+                    self.esm_model.eval()
+                    # Wrap each layer in FSDP separately
+                    for name, child in self.esm_model.named_children():
+                        if name == "layers":
+                            for layer_name, layer in child.named_children():
+                                wrapped_layer = wrap(layer)
+                                setattr(child, layer_name, wrapped_layer)
+                    self.esm_model = wrap(self.esm_model)
+                self.af2_to_esm = self.af2_to_esm.to(self.device)
+            else:
+                self.esm_model = self.esm_model.to(self.device)
+                self.af2_to_esm = self.af2_to_esm.to(self.device)
+
+        if stage == "fit":
+            self.training_gpus = self.trainer.world_size
+            self.hparams["training_gpus"] = self.training_gpus
+
+            batch = next(iter(self.trainer.datamodule.train_dataloader()))
+
+            batch = self.processor.preprocess_training(
+                batch,
+                self.esm_model,
+                self.esm_dict,
+                self.af2_to_esm,
+            )
+            y = batch["coords"]
+            t = torch.zeros((y.shape[0])).cuda()
+
+    def on_train_batch_end(self, outputs, batch, batch_idx):
+        optimizer = self.optimizers()
+        self.log(
+            "trainer/lr",
+            optimizer.param_groups[0]["lr"],
+            on_epoch=True,
+            logger=True,
+            prog_bar=True,
+            rank_zero_only=True,
+        )
+
+    def on_before_optimizer_step(self, optimizer: Optimizer) -> None:
+
+        if isinstance(
+            self.trainer.strategy, lightning.pytorch.strategies.fsdp.FSDPStrategy
+        ):
+
+            with FullyShardedDataParallel.summon_full_params(
+                self.trainer.strategy.model, with_grads=True
+            ):
+                clip_grad_norm_(
+                    self.trainer.strategy.model.model.parameters(),
+                    self.hparams.clip_grad_norm_val,
+                    norm_type=2.0,
+                    error_if_nonfinite=True,
+                )
+
+        else:
+            clip_grad_norm_(
+                self.model.parameters(),
+                self.hparams.clip_grad_norm_val,
+                norm_type=2.0,
+                error_if_nonfinite=True,
+            )
+        return
+
+    def on_before_zero_grad(self, optimizer: Optimizer) -> None:
+        # if self.eval_ema:
+        if isinstance(
+            self.trainer.strategy, lightning.pytorch.strategies.fsdp.FSDPStrategy
+        ):
+
+            with FullyShardedDataParallel.summon_full_params(
+                self.trainer.strategy.model
+            ):
+                self.trainer.strategy.model.model_ema.update_parameters(
+                    self.trainer.strategy.model.model
+                )
+
+        else:
+            self.model_ema.update_parameters(self.model)
+        return
+
+    def on_save_checkpoint(self, checkpoint) -> None:
+
+        if not isinstance(
+            self.trainer.strategy, lightning.pytorch.strategies.fsdp.FSDPStrategy
+        ):
+            layers_to_delete = []
+            for k in checkpoint["state_dict"].keys():
+                if k.startswith("esm_model._fsdp_wrapped_module"):
+                    layers_to_delete.append(k)
+
+            for k in layers_to_delete:
+                del checkpoint["state_dict"][k]
+
+        else:
+            try:
+                del checkpoint["hyper_parameters"]
+            except:
+                print("No hyper_parameters in checkpoint")
+
+        return super().on_save_checkpoint(checkpoint)
+
+    def on_load_checkpoint(self, checkpoint) -> None:
+
+        if not isinstance(
+            self.trainer.strategy, lightning.pytorch.strategies.fsdp.FSDPStrategy
+        ):
+            self.training_gpus = checkpoint["hyper_parameters"]["training_gpus"]
+            self.fwd_flops = checkpoint["hyper_parameters"]["fwd_flops"]
+
+        if checkpoint["loops"] is not None:
+
+            self.trainer.fit_loop.load_state_dict(checkpoint["loops"]["fit_loop"])
+            self.trainer.validate_loop.load_state_dict(
+                checkpoint["loops"]["validate_loop"]
+            )
+        return super().on_load_checkpoint(checkpoint)
+
+    def configure_optimizers(self):
+        optimizer = self.hparams.optimizer(
+            params=self.trainer.model.parameters()
+        )
+
+        if self.hparams.scheduler is not None:
+            scheduler = self.hparams.scheduler(optimizer=optimizer)
+            return {
+                "optimizer": optimizer,
+                "lr_scheduler": {
+                    "scheduler": scheduler,
+                    "interval": "step",
+                    "frequency": 1,
+                },
+            }
+        return {"optimizer": optimizer}

models/simplefold/torch/__init__.py

@@ -0,0 +1,4 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+#

models/simplefold/torch/architecture.py

@@ -0,0 +1,304 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+#
+
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from .layers import FinalLayer, ConditionEmbedder
+from onescience.utils.simplefold.esm_utils import esm_model_dict
+
+
+class FoldingDiT(nn.Module):
+    def __init__(
+        self,
+        trunk,
+        time_embedder,
+        aminoacid_pos_embedder,
+        pos_embedder,
+        atom_encoder_transformer,
+        atom_decoder_transformer,
+        hidden_size=1152,
+        num_heads=16,
+        atom_num_heads=4,
+        output_channels=3,
+        atom_hidden_size_enc=256,
+        atom_hidden_size_dec=256,
+        atom_n_queries_enc=32,
+        atom_n_keys_enc=128,
+        atom_n_queries_dec=32,
+        atom_n_keys_dec=128,
+        esm_model="esm2_3B",
+        esm_dropout_prob=0.0,
+        use_atom_mask=False,
+        use_length_condition=True,
+    ):
+        super().__init__()
+        self.pos_embedder = pos_embedder
+        pos_embed_channels = pos_embedder.embed_dim
+        self.aminoacid_pos_embedder = aminoacid_pos_embedder
+        aminoacid_pos_embed_channels = aminoacid_pos_embedder.embed_dim
+
+        self.time_embedder = time_embedder
+
+        self.atom_encoder_transformer = atom_encoder_transformer
+        self.atom_decoder_transformer = atom_decoder_transformer
+
+        self.trunk = trunk
+
+        self.hidden_size = hidden_size
+        self.output_channels = output_channels
+        self.num_heads = num_heads
+        self.atom_num_heads = atom_num_heads
+        self.use_atom_mask = use_atom_mask
+        self.esm_dropout_prob = esm_dropout_prob
+        self.use_length_condition = use_length_condition
+
+        esm_s_dim = esm_model_dict[esm_model]["esm_s_dim"]
+        esm_num_layers = esm_model_dict[esm_model]["esm_num_layers"]
+
+        self.atom_hidden_size_enc = atom_hidden_size_enc
+        self.atom_hidden_size_dec = atom_hidden_size_dec
+        self.atom_n_queries_enc = atom_n_queries_enc
+        self.atom_n_keys_enc = atom_n_keys_enc
+        self.atom_n_queries_dec = atom_n_queries_dec
+        self.atom_n_keys_dec = atom_n_keys_dec
+
+        atom_feat_dim = pos_embed_channels + aminoacid_pos_embed_channels + 427
+        self.atom_feat_proj = nn.Sequential(
+            nn.Linear(atom_feat_dim, hidden_size),
+            nn.LayerNorm(hidden_size),
+            nn.SiLU(),
+        )
+        self.atom_pos_proj = nn.Linear(pos_embed_channels, hidden_size, bias=False)
+
+        if self.use_length_condition:
+            self.length_embedder = nn.Sequential(
+                nn.Linear(1, hidden_size, bias=False),
+                nn.LayerNorm(hidden_size),
+            )
+
+        self.atom_in_proj = nn.Linear(hidden_size * 2, hidden_size, bias=False)
+
+        self.esm_s_combine = nn.Parameter(torch.zeros(esm_num_layers))
+        self.esm_s_proj = ConditionEmbedder(
+            input_dim=esm_s_dim,
+            hidden_size=hidden_size,
+            dropout_prob=self.esm_dropout_prob,
+        )
+        latent_cat_dim = hidden_size * 2
+        self.esm_cat_proj = nn.Linear(latent_cat_dim, hidden_size)
+
+        self.context2atom_proj = nn.Sequential(
+            nn.Linear(hidden_size, self.atom_hidden_size_enc),
+            nn.LayerNorm(self.atom_hidden_size_enc),
+        )
+        self.atom2latent_proj = nn.Sequential(
+            nn.Linear(self.atom_hidden_size_enc, hidden_size),
+            nn.LayerNorm(hidden_size),
+        )
+        self.atom_enc_cond_proj = nn.Sequential(
+            nn.Linear(hidden_size, self.atom_hidden_size_enc),
+            nn.LayerNorm(self.atom_hidden_size_enc),
+        )
+        self.atom_dec_cond_proj = nn.Sequential(
+            nn.Linear(hidden_size, self.atom_hidden_size_dec),
+            nn.LayerNorm(self.atom_hidden_size_dec),
+        )
+
+        self.latent2atom_proj = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size),
+            nn.SiLU(),
+            nn.LayerNorm(hidden_size),
+            nn.Linear(hidden_size, self.atom_hidden_size_dec),
+        )
+
+        self.final_layer = FinalLayer(
+            self.atom_hidden_size_dec, 
+            output_channels, 
+            c_dim=hidden_size
+        )
+
+    def create_local_attn_bias(
+        self, n: int, n_queries: int, n_keys: int, inf: float = 1e10, device: torch.device = None
+    ) -> torch.Tensor:
+        """Create local attention bias based on query window n_queries and kv window n_keys.
+
+        Args:
+            n (int): the length of quiries
+            n_queries (int): window size of quiries
+            n_keys (int): window size of keys/values
+            inf (float, optional): the inf to mask attention. Defaults to 1e10.
+            device (torch.device, optional): cuda|cpu|None. Defaults to None.
+
+        Returns:
+            torch.Tensor: the diagonal-like global attention bias
+        """
+        n_trunks = int(math.ceil(n / n_queries))
+        padded_n = n_trunks * n_queries
+        attn_mask = torch.zeros(padded_n, padded_n, device=device)
+        for block_index in range(0, n_trunks):
+            i = block_index * n_queries
+            j1 = max(0, n_queries * block_index - (n_keys - n_queries) // 2)
+            j2 = n_queries * block_index + (n_queries + n_keys) // 2
+            attn_mask[i : i + n_queries, j1:j2] = 1.0
+        attn_bias = (1 - attn_mask) * -inf
+        return attn_bias.to(device=device)[:n, :n]
+
+    def create_atom_attn_mask(
+        self, 
+        feats, 
+        natoms, 
+        atom_n_queries=None, 
+        atom_n_keys=None,
+        inf: float = 1e10
+    ) -> torch.Tensor:
+        if atom_n_queries is not None and atom_n_keys is not None:
+            atom_attn_mask = self.create_local_attn_bias(
+                n=natoms,
+                n_queries=atom_n_queries,
+                n_keys=atom_n_keys,
+                device=feats["ref_pos"].device,
+                inf=inf,
+            )
+        else:
+            atom_attn_mask = None
+
+        return atom_attn_mask
+
+    def forward(self, noised_pos, t, feats, self_cond=None):
+        B, N, _ = feats["ref_pos"].shape
+        M = feats["mol_type"].shape[1]
+        atom_to_token = feats["atom_to_token"].float() # [B, N, M]
+        atom_to_token_idx = feats["atom_to_token_idx"]
+        ref_space_uid = feats["ref_space_uid"]
+
+        # create atom attention masks
+        atom_attn_mask_enc = self.create_atom_attn_mask(
+            feats, 
+            natoms=N,
+            atom_n_queries=self.atom_n_queries_enc,
+            atom_n_keys=self.atom_n_keys_enc,
+        )
+        atom_attn_mask_dec = self.create_atom_attn_mask(
+            feats,
+            natoms=N,
+            atom_n_queries=self.atom_n_queries_dec,
+            atom_n_keys=self.atom_n_keys_dec,
+        )
+
+        # create condition embeddings for AdaLN
+        c_emb = self.time_embedder(t)  # (B, D)
+        if self.use_length_condition:
+            length = feats["max_num_tokens"].float().unsqueeze(-1)
+            c_emb = c_emb + self.length_embedder(torch.log(length))
+
+        # create atom features
+        mol_type = feats["mol_type"]
+        mol_type = F.one_hot(mol_type, num_classes=4).float()       # [B, M, 4]
+        res_type = feats["res_type"].float()                        # [B, M, 33]
+        pocket_feature = feats["pocket_feature"].float()            # [B, M, 4]
+        res_feat = torch.cat(
+            [mol_type, res_type, pocket_feature], 
+        dim=-1)                                                     # [B, M, 41]
+        atom_feat_from_res = torch.bmm(atom_to_token, res_feat)     # [B, N, 41]
+        atom_res_pos = self.aminoacid_pos_embedder(
+            pos=atom_to_token_idx.unsqueeze(-1).float()
+        )
+        ref_pos_emb = self.pos_embedder(pos=feats["ref_pos"])
+        atom_feat = torch.cat(
+            [
+                ref_pos_emb,                                        # (B, N, PD1)
+                atom_feat_from_res,                                 # (B, N, 41)
+                atom_res_pos,                                       # (B, N, PD2)
+                feats["ref_charge"].unsqueeze(-1),                  # (B, N, 1)
+                feats["atom_pad_mask"].unsqueeze(-1),               # (B, N, 1)
+                feats["ref_element"],                               # (B, N, 128)
+                feats["ref_atom_name_chars"].reshape(B, N, 4 * 64), # (B, N, 256)
+            ],
+            dim=-1,
+        )                                                           # (B, N, PD1+PD2+427)
+        atom_feat = self.atom_feat_proj(atom_feat)                  # (B, N, D)
+
+        atom_coord = self.pos_embedder(pos=noised_pos)              # (B, N, PD1)
+        atom_coord = self.atom_pos_proj(atom_coord)                 # (B, N, D)
+
+        atom_in = torch.cat([atom_feat, atom_coord], dim=-1)
+        atom_in = self.atom_in_proj(atom_in)                        # (B, N, D)
+
+        # position embeddings for Axial RoPE
+        atom_pe_pos = torch.cat(
+            [
+                ref_space_uid.unsqueeze(-1).float(),                 # (B, N, 1)
+                feats["ref_pos"],                                    # (B, N, 3)
+            ],
+            dim=-1,
+        )                                                            # (B, N, 4)
+        token_pe_pos = torch.cat(
+            [
+                feats["residue_index"].unsqueeze(-1).float(),        # (B, M, 1)
+                feats["entity_id"].unsqueeze(-1).float(),            # (B, M, 1)
+                feats["asym_id"].unsqueeze(-1).float(),              # (B, M, 1)
+                feats["sym_id"].unsqueeze(-1).float(),               # (B, M, 1)
+            ],
+            dim=-1,
+        )                                                            # (B, M, 4)
+
+        # atom encoder
+        atom_c_emb_enc = self.atom_enc_cond_proj(c_emb)
+        atom_latent = self.context2atom_proj(atom_in)
+        atom_latent = self.atom_encoder_transformer(
+            latents=atom_latent, 
+            c=atom_c_emb_enc, 
+            attention_mask=atom_attn_mask_enc,
+            pos=atom_pe_pos,
+        )
+        atom_latent = self.atom2latent_proj(atom_latent)
+
+        # grouping: aggregate atom tokens to residue tokens
+        atom_to_token_mean = atom_to_token / (
+            atom_to_token.sum(dim=1, keepdim=True) + 1e-6
+        )
+        latent = torch.bmm(atom_to_token_mean.transpose(1, 2), atom_latent)
+        assert latent.shape[1] == M
+
+        esm_s = (self.esm_s_combine.softmax(0).unsqueeze(0) @ feats['esm_s']).squeeze(2)
+        force_drop_ids = feats.get("force_drop_ids", None)
+        esm_emb = self.esm_s_proj(esm_s, self.training, force_drop_ids)
+        assert esm_emb.shape[1] == latent.shape[1]
+
+        latent = self.esm_cat_proj(torch.cat([latent, esm_emb], dim=-1))
+
+        # residue trunk
+        latent = self.trunk(
+            latents=latent, 
+            c=c_emb, 
+            attention_mask=None,
+            pos=token_pe_pos,
+        )
+
+        # ungrouping: broadcast residue tokens to atom tokens
+        output = torch.bmm(atom_to_token, latent)
+        assert output.shape[1] == N
+
+        # add skip connection
+        output = output + atom_latent
+        output = self.latent2atom_proj(output)
+
+        # atom decoder
+        atom_c_emb_dec = self.atom_dec_cond_proj(c_emb)
+        output = self.atom_decoder_transformer(
+            latents=output, 
+            c=atom_c_emb_dec,
+            attention_mask=atom_attn_mask_dec,
+            pos=atom_pe_pos,
+        )
+        output = self.final_layer(output, c=c_emb)
+
+        return {
+            "predict_velocity": output,
+            "latent": latent,
+        }

models/simplefold/torch/layers.py

@@ -0,0 +1,313 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+#
+
+import math
+from einops import rearrange
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+#################################################################################
+#                            Attention Layers                                  #
+#################################################################################
+
+
+class SelfAttentionLayer(nn.Module):
+    def __init__(
+        self,
+        hidden_size,
+        num_heads=8,
+        qkv_bias=False,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+        use_bias=True,
+        qk_norm=True,
+        pos_embedder=None,
+        linear_target: nn.Module = nn.Linear,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = hidden_size // num_heads
+        self.scale = qk_scale or head_dim**-0.5
+
+        self.qkv = linear_target(hidden_size, hidden_size * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = linear_target(hidden_size, hidden_size, bias=use_bias)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        self.q_norm = RMSNorm(head_dim) if qk_norm else nn.Identity()
+        self.k_norm = RMSNorm(head_dim) if qk_norm else nn.Identity()
+
+        self.pos_embedder = pos_embedder
+
+    def forward(self, x, **kwargs):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
+        pos = kwargs.get("pos")
+
+        qkv = rearrange(qkv, "b n t h c -> t b h n c")
+        q, k, v = (
+            qkv[0],
+            qkv[1],
+            qkv[2],
+        )  # make torchscript happy (cannot use tensor as tuple)
+
+        q, k = self.q_norm(q), self.k_norm(k)
+
+        if self.pos_embedder and pos is not None:
+            q, k = self.pos_embedder(q, k, pos)
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class EfficientSelfAttentionLayer(SelfAttentionLayer):
+    """Started from https://github.com/facebookresearch/dinov2/blob/main/dinov2/layers/attention.py"""
+
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, x, **kwargs):
+        B, N, C = x.shape
+        attn_mask = kwargs.get("attention_mask")
+        pos = kwargs.get("pos")
+
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
+        qkv = rearrange(qkv, "b n t h c -> t b h n c")
+        q, k, v = qkv.unbind(0)
+
+        if attn_mask is not None:
+            attn_mask = attn_mask.to(dtype=q.dtype)
+
+        if self.pos_embedder and pos is not None:
+            q, k = self.pos_embedder(q, k, pos)
+
+        q, k = self.q_norm(q), self.k_norm(k)
+        v1 = v.to(dtype = q.dtype)
+        x = nn.functional.scaled_dot_product_attention(q, k, v1, attn_mask=attn_mask)
+
+        x = x.transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+
+        return x
+
+
+#################################################################################
+#                              FeedForward Layer                                #
+#################################################################################
+
+
+class SwiGLUFeedForward(nn.Module):
+    def __init__(self, dim, hidden_dim, multiple_of=256):
+        super().__init__()
+        hidden_dim = int(2 * hidden_dim / 3)
+        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
+
+        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, dim, bias=True)
+        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        torch.nn.init.xavier_uniform_(self.w1.weight)
+        torch.nn.init.xavier_uniform_(self.w2.weight)
+        torch.nn.init.xavier_uniform_(self.w3.weight)
+        if self.w1.bias is not None:
+            torch.nn.init.constant_(self.w1.bias, 0)
+        if self.w2.bias is not None:
+            torch.nn.init.constant_(self.w2.bias, 0)
+        if self.w3.bias is not None:
+            torch.nn.init.constant_(self.w3.bias, 0)
+
+    def forward(self, x):
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+#################################################################################
+#                               Utility Layers                                  #
+#################################################################################
+
+
+class TimestepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations.
+    """
+
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        nn.init.normal_(self.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.mlp[2].weight, std=0.02)
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        """
+        Create sinusoidal timestep embeddings.
+        :param t: a 1-D Tensor of N indices, one per batch element.
+                          These may be fractional.
+        :param dim: the dimension of the output.
+        :param max_period: controls the minimum frequency of the embeddings.
+        :return: an (N, D) Tensor of positional embeddings.
+        """
+        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period)
+            * torch.arange(start=0, end=half, dtype=torch.float32)
+            / half
+        ).to(device=t.device)
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat(
+                [embedding, torch.zeros_like(embedding[:, :1])], dim=-1
+            )
+        return embedding
+
+    def forward(self, t):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+class ConditionEmbedder(nn.Module):
+    """
+    Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance.
+    """
+    def __init__(self, input_dim, hidden_size, dropout_prob):
+        super().__init__()
+        self.proj = nn.Sequential(
+                nn.Linear(input_dim, hidden_size),
+                nn.LayerNorm(hidden_size),
+                nn.SiLU(),
+            )
+        self.dropout_prob = dropout_prob
+        self.null_token = nn.Parameter(torch.randn(input_dim), requires_grad=True)
+
+    def token_drop(self, cond, force_drop_ids=None):
+        """
+        cond: (B, N, D)
+        Drops conditions to enable classifier-free guidance.
+        """
+        if force_drop_ids is None:
+            drop_ids = torch.rand(cond.shape[0], device=cond.device) < self.dropout_prob
+        else:
+            drop_ids = force_drop_ids
+        cond[drop_ids] = self.null_token[None, None, :]
+        return cond
+
+    def forward(self, cond, train, force_drop_ids=None):
+        use_dropout = self.dropout_prob > 0
+        if (train and use_dropout) or (force_drop_ids is not None):
+            cond = self.token_drop(cond, force_drop_ids)
+        embeddings = self.proj(cond)
+        return embeddings
+
+
+class FinalLayer(nn.Module):
+    """
+    The final layer of DiT.
+    """
+
+    def __init__(self, hidden_size, out_channels, c_dim=None):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, out_channels, bias=True)
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(), nn.Linear(c_dim, 2 * hidden_size, bias=True)
+        )
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        # Initialize transformer layers:
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+
+        self.apply(_basic_init)
+
+        # Zero-out output layers:
+        nn.init.constant_(self.adaLN_modulation[-1].weight, 0)
+        nn.init.constant_(self.adaLN_modulation[-1].bias, 0)
+        nn.init.constant_(self.linear.weight, 0)
+        nn.init.constant_(self.linear.bias, 0)
+
+    def forward(self, x, c):
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
+        x = modulate(self.norm_final(x), shift, scale)
+        x = self.linear(x)
+        return x
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, d, p=-1.0, eps=1e-8, bias=False):
+        """
+            Root Mean Square Layer Normalization
+        :param d: model size
+        :param p: partial RMSNorm, valid value [0, 1], default -1.0 (disabled)
+        :param eps:  epsilon value, default 1e-8
+        :param bias: whether use bias term for RMSNorm, disabled by
+            default because RMSNorm doesn't enforce re-centering invariance.
+        """
+        super(RMSNorm, self).__init__()
+
+        self.eps = eps
+        self.d = d
+        self.p = p
+        self.bias = bias
+
+        self.scale = nn.Parameter(torch.ones(d))
+        self.register_parameter("scale", self.scale)
+
+        if self.bias:
+            self.offset = nn.Parameter(torch.zeros(d))
+            self.register_parameter("offset", self.offset)
+
+    def forward(self, x):
+        if self.p < 0.0 or self.p > 1.0:
+            norm_x = x.norm(2, dim=-1, keepdim=True, dtype=x.dtype)
+            d_x = self.d
+        else:
+            partial_size = int(self.d * self.p)
+            partial_x, _ = torch.split(x, [partial_size, self.d - partial_size], dim=-1)
+
+            norm_x = partial_x.norm(2, dim=-1, keepdim=True, dtype=x.dtype)
+            d_x = partial_size
+
+        rms_x = norm_x * d_x ** (-1.0 / 2)
+        x_normed = x / (rms_x + self.eps)
+
+        if self.bias:
+            return self.scale * x_normed + self.offset
+
+        return self.scale * x_normed

models/simplefold/torch/pos_embed.py

@@ -0,0 +1,177 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+#
+
+import math
+import torch
+from torch import nn
+
+
+class AbsolutePositionEncoding(nn.Module):
+    def __init__(self, in_dim, embed_dim, include_input=False):
+        super().__init__()
+        self.in_dim = in_dim
+        self.hidden_dim = embed_dim
+        self.include_input = include_input
+        assert embed_dim % in_dim == 0, "embed_dim must be divisible by in_dim"
+        self.embed_dim = embed_dim + in_dim if include_input else embed_dim
+
+    def forward(self, pos):
+        pos_embs = []
+        for i in range(self.in_dim):
+            pe = self.get_1d_pos_embed(pos[..., i])
+            pos_embs.append(pe)
+        if self.include_input:
+            pos_embs.append(pos)
+        pos_embs = torch.cat(pos_embs, dim=-1)
+        return pos_embs
+
+    def get_1d_pos_embed(self, pos):
+        """
+        https://github.com/facebookresearch/DiT/blob/main/models.py#L303
+        """
+        embed_dim = self.hidden_dim // (self.in_dim * 2)
+        omega = 2 ** torch.linspace(0, math.log(224, 2) - 1, embed_dim).to(pos.device)
+        omega *= torch.pi
+
+        if len(pos.shape) == 1:
+            out = torch.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
+        elif len(pos.shape) == 2:
+            out = torch.einsum("nm,d->nmd", pos, omega)
+
+        emb_sin = torch.sin(out)  # (*, M, D/2)
+        emb_cos = torch.cos(out)  # (*, M, D/2)
+        emb = torch.cat([emb_sin, emb_cos], dim=-1)  # (*, M, D)
+        return emb
+
+
+class FourierPositionEncoding(torch.nn.Module):
+    def __init__(
+        self,
+        in_dim: int,
+        include_input: bool = False,
+        min_freq_log2: float = 0,
+        max_freq_log2: float = 12,
+        num_freqs: int = 32,
+        log_sampling: bool = True,
+    ):
+        super().__init__()
+        self.in_dim = in_dim
+        self.include_input = include_input
+        self.min_freq_log2 = min_freq_log2
+        self.max_freq_log2 = max_freq_log2
+        self.num_freqs = num_freqs
+        self.log_sampling = log_sampling
+        self.create_embedding_fn()
+
+    def create_embedding_fn(self):
+        d = self.in_dim
+        dim_out = 0
+        if self.include_input:
+            dim_out += d
+
+        min_freq = self.min_freq_log2
+        max_freq = self.max_freq_log2
+        N_freqs = self.num_freqs
+
+        if self.log_sampling:
+            freq_bands = 2.0 ** torch.linspace(
+                min_freq, max_freq, steps=N_freqs
+            )  # (nf,)
+        else:
+            freq_bands = torch.linspace(
+                2.0**min_freq, 2.0**max_freq, steps=N_freqs
+            )  # (nf,)
+
+        assert (
+            freq_bands.isfinite().all()
+        ), f"nan: {freq_bands.isnan().any()} inf: {freq_bands.isinf().any()}"
+
+        self.register_buffer("freq_bands", freq_bands)  # (nf,)
+        self.embed_dim = dim_out + d * self.freq_bands.numel() * 2
+
+    def forward(
+        self,
+        pos: torch.Tensor,
+    ):
+        """
+        Get the positional encoding for each coordinate.
+        Args:
+            pos:
+                (*, in_dim)
+        Returns:
+            out:
+                (*, in_dimitional_encoding)
+        """
+
+        out = []
+        if self.include_input:
+            out = [pos]  # (*, in_dim)
+
+        pos = pos.unsqueeze(-1) * self.freq_bands  # (*b, d, nf)
+
+        out += [
+            torch.sin(pos).flatten(start_dim=-2),  # (*b, d*nf)
+            torch.cos(pos).flatten(start_dim=-2),  # (*b, d*nf)
+        ]
+
+        out = torch.cat(out, dim=-1)  # (*b, 2 * in_dim * nf (+ in_dim))
+        return out
+
+
+def compute_axial_cis(
+    ts: torch.Tensor,
+    in_dim: int,
+    dim: int,
+    theta: float = 100.0,
+):
+    B, N, D = ts.shape
+    freqs_all = []
+    interval = 2 * in_dim
+    for i in range(in_dim):
+        freq = 1.0 / (
+            theta ** (torch.arange(0, dim, interval)[: (dim // interval)].float() / dim)
+        ).to(ts.device)
+        t = ts[..., i].flatten()
+        freq_i = torch.outer(t, freq)
+        freq_cis_i = torch.polar(torch.ones_like(freq_i), freq_i)
+        freq_cis_i = freq_cis_i.view(B, N, -1)
+        freqs_all.append(freq_cis_i)
+    freqs_cis = torch.cat(freqs_all, dim=-1)
+    return freqs_cis
+
+
+def apply_rotary_emb(xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor):
+    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
+    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
+    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
+    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
+    return xq_out.type_as(xq).to(xq.device), xk_out.type_as(xk).to(xk.device)
+
+
+class AxialRotaryPositionEncoding(nn.Module):
+    def __init__(
+        self,
+        in_dim,
+        embed_dim,
+        num_heads,
+        base=100.0,
+    ):
+        super().__init__()
+        self.in_dim = in_dim
+        self.num_heads = num_heads
+        self.embed_dim = embed_dim // num_heads
+        self.base = base
+
+    def forward(self, xq, xk, pos):
+        """
+        xq: [B, H, N, D]
+        xk: [B, H, N, D]
+        pos: [B, N, in_dim]
+        """
+        if pos.ndim == 2:
+            pos = pos.unsqueeze(-1)
+        freqs_cis = compute_axial_cis(pos, self.in_dim, self.embed_dim, self.base)
+        freqs_cis = freqs_cis.unsqueeze(1)
+        return apply_rotary_emb(xq, xk, freqs_cis.to(xq.device))

scripts/evaluate.py

@@ -0,0 +1,101 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+#
+
+import hydra
+import torch
+from omegaconf import OmegaConf
+import lightning.pytorch as pl
+import sys
+from pathlib import Path
+from lightning.pytorch import (
+    LightningDataModule, 
+    LightningModule
+)
+
+ROOT = Path(__file__).resolve().parents[1]
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+from onescience.utils.simplefold.utils import extras, create_folders, task_wrapper
+from onescience.utils.simplefold.instantiators import instantiate_callbacks
+from onescience.utils.simplefold.logging_utils import log_hyperparameters
+from onescience.utils.simplefold.pylogger import RankedLogger
+
+torch.set_float32_matmul_precision("medium")
+log = RankedLogger(__name__, rank_zero_only=True)
+
+
+@task_wrapper
+def test(cfg):
+    load_ckpt_path = cfg.get("load_ckpt_path", None)
+    assert load_ckpt_path != None
+
+    log.info(f"Instantiating model <{cfg.model._target_}>")
+    model: LightningModule = hydra.utils.instantiate(cfg.model)
+
+    checkpoint = torch.load(load_ckpt_path, map_location="cpu", weights_only=False)
+
+    # TODO: add load FSDP checkpoint
+    try:
+        # checkpoint in our official release only contains weights of EMA model
+        # therefore, we by default load EMA model weight to LightningModule for inference
+        for key in model.model_ema.state_dict().keys():
+            if key.startswith("module."):
+                model.model_ema.state_dict()[key].copy_(
+                    checkpoint[key.replace("module.", "")]
+                )
+        print("Loaded weights of EMA model successfully.")
+    except:
+        # if using checkpoint from your own training, load weights of LightningModule directly
+        model.load_state_dict(
+            checkpoint["state_dict"], strict=False
+        )
+        print("Loaded weights of LightningModule successfully.")
+
+    # reset ESM model to avoid issues in loading FSDP checkpoint
+    model.reset_esm(cfg.model.esm_model)
+
+    seed = cfg.get("seed", 42)
+    pl.seed_everything(seed, workers=True)
+
+    log.info(f"Instantiating datamodule <{cfg.data._target_}>")
+    datamodule: LightningDataModule = hydra.utils.instantiate(cfg.data)
+
+    log.info("Instantiating callbacks...")
+    callbacks = instantiate_callbacks(cfg.get("callbacks"))
+
+    log.info(f"Instantiating trainer <{cfg.trainer._target_}>")
+    trainer = hydra.utils.instantiate(
+        cfg.trainer, callbacks=callbacks, logger=[], plugins=[]
+    )
+
+    object_dict = {
+        "cfg": cfg,
+        "datamodule": datamodule,
+        "model": model,
+        "callbacks": callbacks,
+        "logger": [],
+        "trainer": trainer,
+    }
+
+    if log:
+        log.info("Logging hyperparameters!")
+        log_hyperparameters(object_dict)
+
+    log.info("Starting evaluation!")
+    trainer.predict(model=model, datamodule=datamodule, ckpt_path=None)
+
+
+@hydra.main(version_base="1.3", config_path="../config", config_name="base_eval.yaml")
+def submit_run(cfg):
+    OmegaConf.resolve(cfg)
+    extras(cfg)
+    create_folders(cfg)
+    test(cfg)
+    return
+
+
+if __name__ == "__main__":
+    submit_run()

scripts/inference.py

@@ -0,0 +1,420 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+#
+
+import os
+import sys
+import torch
+import hydra
+import omegaconf
+import argparse
+import numpy as np
+from copy import deepcopy
+from pathlib import Path
+from itertools import starmap
+import lightning.pytorch as pl
+
+ROOT = Path(__file__).resolve().parents[1]
+CONFIG_DIR = ROOT / "config"
+WEIGHT_DIR = ROOT / "weight"
+
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+from models.simplefold.flow import LinearPath
+from models.simplefold.torch.sampler import EMSampler
+
+from onescience.datapipes.simplefold.processor.protein_processor import ProteinDataProcessor
+from onescience.utils.simplefold.datamodule_utils import process_one_inference_structure
+from onescience.utils.simplefold.esm_utils import _af2_to_esm, esm_registry
+import models.esm.pretrained as esm_pretrained
+from onescience.utils.simplefold.boltz_utils import process_structure, save_structure
+from onescience.utils.simplefold.fasta_utils import process_fastas, check_fasta_inputs
+from onescience.datapipes.boltz_data_pipeline.feature.featurizer import BoltzFeaturizer
+from onescience.datapipes.boltz_data_pipeline.tokenize.boltz_protein import BoltzTokenizer
+
+try: 
+    import mlx.core as mx
+    from mlx.utils import tree_unflatten, tree_flatten
+    from models.simplefold.mlx.sampler import EMSampler as EMSamplerMLX
+    from models.simplefold.mlx.esm_network import ESM2 as ESM2MLX
+    from onescience.utils.simplefold.mlx_utils import map_torch_to_mlx, map_plddt_torch_to_mlx
+    MLX_AVAILABLE = True
+except:
+    MLX_AVAILABLE = False
+    print("MLX not installed, skip importing MLX related packages.")
+
+
+SUPPORTED_MODELS = {
+    "simplefold_100M",
+    "simplefold_360M",
+    "simplefold_700M",
+    "simplefold_1.1B",
+    "simplefold_1.6B",
+    "simplefold_3B",
+}
+
+ESM2_MODEL_NAME = "esm2_t36_3B_UR50D"
+MIN_REAL_WEIGHT_BYTES = 1024
+
+
+def _resolve_path(path: str | os.PathLike[str], base: Path = ROOT) -> Path:
+    path = Path(path)
+    return path if path.is_absolute() else base / path
+
+
+def _require_local_file(path: Path, description: str, min_bytes: int = MIN_REAL_WEIGHT_BYTES) -> Path:
+    if not path.exists():
+        raise FileNotFoundError(
+            f"Missing {description}: {path}\n"
+            f"Please place the file under {WEIGHT_DIR} or pass the corresponding environment/CLI path."
+        )
+    if path.stat().st_size < min_bytes:
+        raise RuntimeError(
+            f"{description} looks like a link/placeholder rather than a real weight file: {path}\n"
+            "Replace it with the real file before running inference."
+        )
+    return path
+
+
+def load_local_or_hub_esm2_3b():
+    candidate_paths = []
+
+    esm_model_path = os.getenv("SIMPLEFOLD_ESM2_MODEL_PATH")
+    if esm_model_path:
+        candidate_paths.append(esm_model_path)
+
+    candidate_paths.extend(
+        [
+            WEIGHT_DIR / "esm_models" / f"{ESM2_MODEL_NAME}.pt",
+            WEIGHT_DIR / f"{ESM2_MODEL_NAME}.pt",
+        ]
+    )
+
+    seen_paths = set()
+    for candidate_path in candidate_paths:
+        if candidate_path in seen_paths:
+            continue
+        seen_paths.add(candidate_path)
+
+        candidate_path = Path(candidate_path)
+        regression_path = candidate_path.with_name(
+            f"{candidate_path.stem}-contact-regression.pt"
+        )
+        if candidate_path.exists() and regression_path.exists():
+            _require_local_file(candidate_path, "ESM-2 3B model weight")
+            _require_local_file(regression_path, "ESM-2 contact regression weight")
+            print(f"Loading ESM-2 3B weights from local path: {candidate_path}")
+            return esm_pretrained.load_model_and_alphabet(str(candidate_path))
+
+    raise FileNotFoundError(
+        "Local ESM-2 3B weights were not found. Put both files here:\n"
+        f"  {WEIGHT_DIR / 'esm_models' / f'{ESM2_MODEL_NAME}.pt'}\n"
+        f"  {WEIGHT_DIR / 'esm_models' / f'{ESM2_MODEL_NAME}-contact-regression.pt'}\n"
+        "or set SIMPLEFOLD_ESM2_MODEL_PATH to the .pt file. Remote download is disabled in this package."
+    )
+
+
+def resolve_ccd_path(cache: Path) -> Path:
+    candidate_paths = []
+
+    ccd_path = os.getenv("SIMPLEFOLD_CCD_PATH")
+    if ccd_path:
+        candidate_paths.append(Path(ccd_path))
+
+    candidate_paths.append(WEIGHT_DIR / "ccd.pkl")
+
+    seen_paths = set()
+    for candidate_path in candidate_paths:
+        resolved_path = str(candidate_path)
+        if resolved_path in seen_paths:
+            continue
+        seen_paths.add(resolved_path)
+
+        if candidate_path.exists():
+            _require_local_file(candidate_path, "CCD dictionary")
+            print(f"Using local CCD dictionary: {candidate_path}")
+            return candidate_path
+
+    raise FileNotFoundError(
+        f"Missing CCD dictionary. Put ccd.pkl under {WEIGHT_DIR} or set SIMPLEFOLD_CCD_PATH."
+    )
+
+
+def get_config_path(relative_path):
+    """Resolve a SimpleFold config from the bundled config directory."""
+    config_subpath = relative_path.replace("configs/", "")
+    config_path = CONFIG_DIR / config_subpath
+    if not config_path.is_file():
+        raise FileNotFoundError(f"Could not find config file: {config_path}")
+    return str(config_path)
+
+
+
+def initialize_folding_model(args):
+    # define folding model
+    simplefold_model = args.simplefold_model
+    if simplefold_model not in SUPPORTED_MODELS:
+        raise ValueError(f"Unsupported model {simplefold_model!r}. Choose one of {sorted(SUPPORTED_MODELS)}")
+
+    # create checkpoint directory
+    ckpt_dir = _resolve_path(args.ckpt_dir)
+    ckpt_path = os.path.join(ckpt_dir, f"{simplefold_model}.ckpt")
+
+    # create folding model
+    ckpt_path = os.path.join(ckpt_dir, f"{simplefold_model}.ckpt")
+    _require_local_file(Path(ckpt_path), f"{simplefold_model} checkpoint")
+    cfg_path = get_config_path(f"configs/model/architecture/foldingdit_{simplefold_model[11:]}.yaml")
+
+    checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+
+    # load model checkpoint
+    if args.backend == 'torch':
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        model_config = omegaconf.OmegaConf.load(cfg_path)
+        model = hydra.utils.instantiate(model_config)
+        model.load_state_dict(checkpoint, strict=True)
+        model = model.to(device)
+    elif args.backend == 'mlx':
+        device = "cpu"
+        # replace torch implementations with mlx
+        with open(cfg_path, "r") as f:
+            yaml_str = f.read()
+        yaml_str = yaml_str.replace('torch', 'mlx')
+
+        model_config = omegaconf.OmegaConf.create(yaml_str)
+        model = hydra.utils.instantiate(model_config)
+        mlx_state_dict = {k: mx.array(v) for k, v in starmap(map_torch_to_mlx, checkpoint.items()) if k is not None}
+        model.update(tree_unflatten(list(mlx_state_dict.items())))
+    print(f"Folding model {simplefold_model} loaded.")
+    print(f"Using device: {device}.")
+
+    model.eval()
+    return model, device
+
+
+def initialize_plddt_module(args, device):
+    if not args.plddt:
+        return None, None
+
+    # load pLDDT module if specified
+    ckpt_dir = _resolve_path(args.ckpt_dir)
+    plddt_ckpt_path = ckpt_dir / "plddt.ckpt"
+    if not plddt_ckpt_path.exists():
+        plddt_ckpt_path = ckpt_dir / "plddt_module_1.6B.ckpt"
+    _require_local_file(plddt_ckpt_path, "pLDDT checkpoint")
+
+    plddt_module_path = get_config_path("configs/model/architecture/plddt_module.yaml")
+    plddt_checkpoint = torch.load(plddt_ckpt_path, map_location="cpu", weights_only=False)
+
+    if args.backend == "torch":
+        plddt_config = omegaconf.OmegaConf.load(plddt_module_path)
+        plddt_out_module = hydra.utils.instantiate(plddt_config)
+        plddt_out_module.load_state_dict(plddt_checkpoint, strict=True)
+        plddt_out_module = plddt_out_module.to(device)
+    elif args.backend == "mlx":
+        # replace torch implementations with mlx
+        with open(plddt_module_path, "r") as f:
+            yaml_str = f.read()
+        yaml_str = yaml_str.replace('torch', 'mlx')
+
+        plddt_config = omegaconf.OmegaConf.create(yaml_str)
+        plddt_out_module = hydra.utils.instantiate(plddt_config)
+
+        mlx_state_dict = {k: mx.array(v) for k, v in starmap(map_plddt_torch_to_mlx, plddt_checkpoint.items()) if k is not None}
+        plddt_out_module.update(tree_unflatten(list(mlx_state_dict.items())))
+
+    plddt_out_module.eval()
+    print(f"pLDDT output module loaded with {args.backend} backend.")
+
+    plddt_latent_ckpt_path = ckpt_dir / "simplefold_1.6B.ckpt"
+    _require_local_file(plddt_latent_ckpt_path, "SimpleFold 1.6B pLDDT latent checkpoint")
+
+    plddt_latent_config_path = get_config_path("configs/model/architecture/foldingdit_1.6B.yaml")
+    plddt_latent_checkpoint = torch.load(plddt_latent_ckpt_path, map_location="cpu", weights_only=False)
+
+    if args.backend == "torch":
+        plddt_latent_config = omegaconf.OmegaConf.load(plddt_latent_config_path)
+        plddt_latent_module = hydra.utils.instantiate(plddt_latent_config)
+        plddt_latent_module.load_state_dict(plddt_latent_checkpoint, strict=True)
+        plddt_latent_module = plddt_latent_module.to(device)
+    elif args.backend == "mlx":
+        # replace torch implementations with mlx
+        with open(plddt_latent_config_path, "r") as f:
+            yaml_str = f.read()
+        yaml_str = yaml_str.replace('torch', 'mlx')
+
+        plddt_latent_config = omegaconf.OmegaConf.create(yaml_str)
+        plddt_latent_module = hydra.utils.instantiate(plddt_latent_config)
+        mlx_state_dict = {k: mx.array(v) for k, v in starmap(map_torch_to_mlx, plddt_latent_checkpoint.items()) if k is not None}
+        plddt_latent_module.update(tree_unflatten(list(mlx_state_dict.items())))
+
+    plddt_latent_module.eval()
+    print(f"pLDDT latent module loaded with {args.backend} backend.")
+
+    return plddt_latent_module, plddt_out_module
+
+
+def initialize_esm_model(args, device):
+    # load ESM2 model
+    esm_model, esm_dict = load_local_or_hub_esm2_3b()
+    af2_to_esm = _af2_to_esm(esm_dict)
+
+    if args.backend == 'torch':
+        esm_model = esm_model.to(device)
+        af2_to_esm = af2_to_esm.to(device)
+    elif args.backend == 'mlx':
+        esm_model_mlx = ESM2MLX(num_layers=36, embed_dim=2560, attention_heads=40)
+        esm_state_dict_torch = esm_model.cpu().state_dict()
+
+        esm_state_dict_torch = {k: mx.array(v) for k, v in starmap(map_torch_to_mlx, esm_state_dict_torch.items()) if k is not None}
+        esm_model_mlx.update(tree_unflatten(list(esm_state_dict_torch.items())))
+        esm_model = esm_model_mlx
+    print(f"pLM ESM-3B loaded with {args.backend} backend.")
+
+    esm_model.eval()
+    return esm_model, esm_dict, af2_to_esm
+
+
+def initialize_others(args, device):
+    # prepare data tokenizer, featurizer, and processor
+    tokenizer = BoltzTokenizer()
+    featurizer = BoltzFeaturizer()
+    processor = ProteinDataProcessor(
+        device=device,
+        scale=16.0, 
+        ref_scale=5.0, 
+        multiplicity=1,
+        inference_multiplicity=args.nsample_per_protein,
+        backend=args.backend,
+    )
+
+    # define flow process and sampler
+    flow = LinearPath()
+
+    if args.backend == "torch":
+        sampler_cls = EMSampler
+    elif args.backend == "mlx":
+        sampler_cls = EMSamplerMLX
+
+    sampler = sampler_cls(
+        num_timesteps=args.num_steps,
+        t_start=1e-4,
+        tau=args.tau,
+        log_timesteps=True,
+        w_cutoff=0.99,
+    )
+    return tokenizer, featurizer, processor, flow, sampler
+
+
+def generate_structure(
+    args, batch, sampler, flow, processor,
+    model, plddt_latent_module, plddt_out_module, device
+):
+    # run inference for target protein
+    if args.backend == "torch":
+        noise = torch.randn_like(batch['coords']).to(device)
+    elif args.backend == "mlx":
+        noise = mx.random.normal(batch['coords'].shape)
+    out_dict = sampler.sample(model, flow, noise, batch)
+
+    if args.plddt:
+        if args.backend == "torch":
+            t = torch.ones(batch['coords'].shape[0], device=device)
+            # use unscaled coords to extract latent for pLDDT prediction
+            out_feat = plddt_latent_module(
+                out_dict["denoised_coords"].detach(), t, batch)
+            plddt_out_dict = plddt_out_module(
+                out_feat["latent"].detach(),
+                batch,
+            )
+        elif args.backend == "mlx":
+            t = mx.ones(batch['coords'].shape[0])
+            # use unscaled coords to extract latent for pLDDT prediction
+            out_feat = plddt_latent_module(
+                out_dict["denoised_coords"], t, batch)
+            plddt_out_dict = plddt_out_module(
+                out_feat["latent"],
+                batch,
+            )
+        # scale pLDDT to [0, 100]
+        plddts = plddt_out_dict["plddt"] * 100.0
+    else:
+        plddts = None
+
+    out_dict = processor.postprocess(out_dict, batch)
+    # sampled_coord = out_dict['denoised_coords'].detach()
+    if args.backend == "torch":
+        sampled_coord = out_dict['denoised_coords'].detach()
+    else:
+        sampled_coord = out_dict['denoised_coords']
+
+    pad_mask = batch['atom_pad_mask']
+    return sampled_coord, pad_mask, plddts
+
+
+def predict_structures_from_fastas(args):
+    # create output directories
+    output_dir = _resolve_path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    prediction_dir = output_dir / f"predictions_{args.simplefold_model}"
+    prediction_dir.mkdir(parents=True, exist_ok=True)
+    cache = output_dir / "cache"
+    cache.mkdir(parents=True, exist_ok=True)
+
+    # set random seed for reproducibility
+    pl.seed_everything(args.seed, workers=True)
+
+    if args.backend == "mlx" and not MLX_AVAILABLE:
+        args.backend = "torch"
+        print("MLX not available, switch to torch backend.")
+
+    # initialize models
+    model, device = initialize_folding_model(args)
+    plddt_latent_module, plddt_out_module = initialize_plddt_module(args, device)
+    esm_model, esm_dict, af2_to_esm = initialize_esm_model(args, device)
+
+    # initialize other components
+    tokenizer, featurizer, processor, flow, sampler = initialize_others(args, device)
+
+    # process fasta files to input format
+    ccd_path = resolve_ccd_path(cache)
+    data = check_fasta_inputs(_resolve_path(args.fasta_path))
+    if not data:
+        raise ValueError("No valid input files found. Please check the input directory.")
+    process_fastas(
+        data=data,
+        out_dir=output_dir,
+        ccd_path=ccd_path,
+    )
+
+    for struct_file in output_dir.glob("structures/*.npz"):
+        record_file = output_dir / "records" / f"{struct_file.stem}.json"
+
+        # prepare the target protein data for inference
+        batch, structure, record = process_one_inference_structure(
+            struct_file, record_file,
+            tokenizer, featurizer, processor,
+            esm_model, esm_dict, af2_to_esm,
+        )
+
+        sampled_coord, pad_mask, plddts = generate_structure(
+            args, batch, sampler, flow, processor,
+            model, plddt_latent_module, plddt_out_module, device
+        )
+
+        for i in range(args.nsample_per_protein):
+            sampled_coord_i = sampled_coord[i]
+            pad_mask_i = pad_mask[i]
+
+            # save the generated structure
+            structure_save = process_structure(
+                deepcopy(structure), sampled_coord_i, pad_mask_i, record, backend=args.backend
+            )
+            outname = f"{record.id}_sampled_{i}"
+            save_structure(
+                structure_save, prediction_dir, outname,
+                output_format=args.output_format,
+                plddts=plddts[i] if plddts is not None else None
+            )

scripts/preflight.py

@@ -0,0 +1,153 @@
+#!/usr/bin/env python3
+"""Preflight checks for the packaged SimpleFold ModelScope layout."""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+REQUIRED_FILES = [
+    "README.md",
+    "configuration.json",
+    "requirements_dcu.txt",
+    "env_install.sh",
+    "config/train.yaml",
+    "config/base_train.yaml",
+    "config/base_eval.yaml",
+    "config/model/architecture/foldingdit_100M.yaml",
+    "config/model/architecture/foldingdit_1.6B.yaml",
+    "config/model/architecture/plddt_module.yaml",
+    "scripts/run_inference.py",
+    "scripts/inference.py",
+    "scripts/train.py",
+    "scripts/train_fsdp.py",
+    "scripts/evaluate.py",
+    "scripts/process_data.py",
+    "scripts/tokenize_data.py",
+    "examples/minimal.fasta",
+    "modules/models/simplefold/flow.py",
+    "modules/datapipes/simplefold/processor/protein_processor.py",
+    "modules/utils/simplefold/esm_utils.py",
+    "modules/datapipes/boltz_data_pipeline/types.py",
+    "modules/models/esm/pretrained.py",
+]
+
+REQUIRED_WEIGHTS = [
+    "simplefold_100M.ckpt",
+    "simplefold_360M.ckpt",
+    "simplefold_700M.ckpt",
+    "simplefold_1.1B.ckpt",
+    "simplefold_1.6B.ckpt",
+    "simplefold_3B.ckpt",
+    "plddt.ckpt",
+    "plddt_module_1.6B.ckpt",
+    "ccd.pkl",
+    "boltz1_conf.ckpt",
+]
+
+OPTIONAL_ESM_WEIGHTS = [
+    "esm_models/esm2_t36_3B_UR50D.pt",
+    "esm_models/esm2_t36_3B_UR50D-contact-regression.pt",
+]
+
+
+def ok(message: str) -> None:
+    print(f"[OK] {message}")
+
+
+def warn(message: str) -> None:
+    print(f"[WARN] {message}")
+
+
+def fail(message: str) -> None:
+    print(f"[FAIL] {message}")
+    raise SystemExit(1)
+
+
+def check_files() -> None:
+    missing = [path for path in REQUIRED_FILES if not (ROOT / path).exists()]
+    if missing:
+        fail("Missing required files: " + ", ".join(missing))
+    ok("Required package files are present")
+
+
+def check_weights(strict: bool) -> None:
+    missing = []
+    placeholders = []
+    for name in REQUIRED_WEIGHTS:
+        path = ROOT / "weight" / name
+        if not path.exists():
+            missing.append(name)
+        elif path.stat().st_size < 1024:
+            placeholders.append(name)
+    if missing:
+        fail("Missing weight files: " + ", ".join(missing))
+    if placeholders:
+        message = (
+            "These files look like link/placeholder files and must be replaced "
+            "before real inference/training: " + ", ".join(placeholders)
+        )
+        if strict:
+            fail(message)
+        warn(message)
+    else:
+        ok("Required SimpleFold weights look like real files")
+
+    missing_esm = [name for name in OPTIONAL_ESM_WEIGHTS if not (ROOT / "weight" / name).exists()]
+    if missing_esm:
+        warn(
+            "ESM-2 3B local weights are not present. Put them under weight/esm_models "
+            "before running inference: " + ", ".join(missing_esm)
+        )
+    else:
+        ok("Local ESM-2 3B files are present")
+
+
+def check_imports(strict: bool) -> None:
+    modules = [
+        "modules.models.simplefold.flow",
+        "modules.models.simplefold.torch.architecture",
+        "modules.datapipes.simplefold.processor.protein_processor",
+        "modules.utils.simplefold.esm_utils",
+        "modules.datapipes.boltz_data_pipeline.types",
+        "modules.models.esm.pretrained",
+    ]
+    missing_dependencies = []
+    for module_name in modules:
+        try:
+            importlib.import_module(module_name)
+        except ModuleNotFoundError as exc:
+            missing_dependencies.append(f"{module_name}: {exc.name}")
+    if missing_dependencies:
+        message = (
+            "Some imports need external Python dependencies. Install requirements_dcu.txt "
+            "and rerun with --strict-imports for a full check: "
+            + "; ".join(missing_dependencies)
+        )
+        if strict:
+            fail(message)
+        warn(message)
+    else:
+        ok("Bundled Python modules import successfully")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Check packaged SimpleFold runtime files.")
+    parser.add_argument("--strict-weights", action="store_true")
+    parser.add_argument("--strict-imports", action="store_true")
+    args = parser.parse_args()
+
+    check_files()
+    check_weights(strict=args.strict_weights)
+    check_imports(strict=args.strict_imports)
+    ok("SimpleFold package preflight finished")
+
+
+if __name__ == "__main__":
+    main()

scripts/process_data.py

@@ -0,0 +1,73 @@
+#!/usr/bin/env python
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+#
+
+import argparse
+import multiprocessing
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+from modules.datapipes.simplefold.process_mmcif import process
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Process MMCIF data.")
+    parser.add_argument(
+        "--data_dir",
+        type=Path,
+        required=True,
+        help="The directory containing the MMCIF files.",
+    )
+    parser.add_argument(
+        "--out_dir",
+        type=Path,
+        default="data",
+        help="The output directory.",
+    )
+    parser.add_argument(
+        "--num-processes",
+        type=int,
+        default=multiprocessing.cpu_count(),
+        help="The number of processes.",
+    )
+    parser.add_argument(
+        "--redis-host",
+        type=str,
+        default="localhost",
+        help="The Redis host, only used with --use-redis.",
+    )
+    parser.add_argument(
+        "--redis-port",
+        type=int,
+        default=7777,
+        help="The Redis port, only used with --use-redis.",
+    )
+    parser.add_argument(
+        "--use-redis",
+        action="store_true",
+        help="Use an external Redis CCD resource instead of local weight/ccd.pkl.",
+    )
+    parser.add_argument(
+        "--ccd-path",
+        type=Path,
+        default=ROOT / "weight" / "ccd.pkl",
+        help="Path to the local CCD pickle file.",
+    )
+    parser.add_argument(
+        "--use-assembly",
+        action="store_true",
+        help="Whether to use assembly 1.",
+    )
+    parser.add_argument(
+        "--max-file-size",
+        type=int,
+        default=None,
+    )
+    args = parser.parse_args()
+    process(args)

scripts/run_inference.py

@@ -0,0 +1,41 @@
+#!/usr/bin/env python3
+"""Run SimpleFold structure prediction from the packaged ModelScope layout."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+from inference import predict_structures_from_fastas
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Run SimpleFold structure prediction.")
+    parser.add_argument("--simplefold_model", default="simplefold_100M")
+    parser.add_argument("--ckpt_dir", default="weight")
+    parser.add_argument("--output_dir", default="outputs/minimal_inference")
+    parser.add_argument("--num_steps", type=int, default=10)
+    parser.add_argument("--tau", type=float, default=0.01)
+    parser.add_argument("--no_log_timesteps", action="store_true")
+    parser.add_argument("--fasta_path", default="examples/minimal.fasta")
+    parser.add_argument("--nsample_per_protein", type=int, default=1)
+    parser.add_argument("--plddt", action="store_true")
+    parser.add_argument("--output_format", default="mmcif", choices=["pdb", "mmcif"])
+    parser.add_argument("--backend", default="torch", choices=["torch", "mlx"])
+    parser.add_argument("--seed", type=int, default=42)
+    return parser
+
+
+def main() -> None:
+    args = build_parser().parse_args()
+    predict_structures_from_fastas(args)
+
+
+if __name__ == "__main__":
+    main()
+

scripts/tokenize_data.py

@@ -0,0 +1,58 @@
+#!/usr/bin/env python
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+#
+
+import argparse
+import sys
+from pathlib import Path
+from tqdm import tqdm
+
+ROOT = Path(__file__).resolve().parents[1]
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+from onescience.datapipes.boltz_data_pipeline.tokenize.boltz_protein import BoltzTokenizer
+from onescience.datapipes.boltz_data_pipeline.types import Manifest
+from onescience.datapipes.simplefold.process_structure import tokenize_structure, finalize
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Tokenize structure data.")
+    parser.add_argument(
+        "--target_dir",
+        type=str,
+        required=True,
+        help="Directory containing the processed structure data.",
+    )
+    parser.add_argument(
+        "--token_dir",
+        type=str,
+        required=True,
+        help="Directory to save the tokenized data.",
+    )
+    args = parser.parse_args()
+
+    target_dir = Path(args.target_dir)
+    manifest_path = target_dir / "manifest.json"
+    manifest: Manifest = Manifest.load(manifest_path)
+    tokenizer = BoltzTokenizer()
+    records = manifest.records
+    print(f"Number of records after filtering: {len(records)}")
+
+    save_token_dir = Path(args.token_dir) / "tokens"
+    save_token_record_dir = Path(args.token_dir) / "records"
+    save_token_dir.mkdir(parents=True, exist_ok=True)
+    save_token_record_dir.mkdir(parents=True, exist_ok=True)
+
+    for record in tqdm(records):
+        tokenize_structure(
+            record,
+            tokenizer,
+            target_dir,
+            str(save_token_dir),
+            save_token_record_dir,
+        )
+
+    finalize(Path(args.token_dir))

scripts/train.py

@@ -0,0 +1,142 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+#
+
+import torch
+import lightning.pytorch as pl
+from lightning.pytorch import LightningDataModule, LightningModule
+import hydra
+from omegaconf import OmegaConf
+import os
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+from onescience.utils.simplefold.utils import (
+    extras,
+    create_folders,
+    task_wrapper,
+)
+from onescience.utils.simplefold.instantiators import (
+    instantiate_callbacks,
+    instantiate_loggers,
+    instantiate_trainer,
+)
+from onescience.utils.simplefold.logging_utils import log_hyperparameters
+from onescience.utils.simplefold.pylogger import RankedLogger
+
+log = RankedLogger(__name__, rank_zero_only=True)
+
+ESM2_MODEL_NAME = "esm2_t36_3B_UR50D"
+
+torch.set_float32_matmul_precision("medium")
+torch.backends.cuda.matmul.allow_tf32 = True # This flag defaults to False
+torch.backends.cudnn.allow_tf32 = True       # This flag defaults to True
+
+
+def configure_local_esm_registry(esm_model: str | None) -> None:
+    if esm_model != "esm2_3B":
+        return
+
+    model_path = Path(
+        os.getenv(
+            "SIMPLEFOLD_ESM2_MODEL_PATH",
+            ROOT / "weight" / "esm_models" / f"{ESM2_MODEL_NAME}.pt",
+        )
+    )
+    regression_path = model_path.with_name(f"{model_path.stem}-contact-regression.pt")
+
+    if not model_path.exists() or not regression_path.exists():
+        raise FileNotFoundError(
+            "Local ESM-2 3B weights were not found. Expected files:\n"
+            f"  {model_path}\n"
+            f"  {regression_path}\n"
+            "Set SIMPLEFOLD_ESM2_MODEL_PATH to the local .pt file if it lives elsewhere."
+        )
+
+    import models.esm.pretrained as local_esm_pretrained
+    import onescience.utils.simplefold.esm_utils as esm_utils
+
+    def load_local_esm2_3b():
+        print(f"Loading ESM-2 3B weights from local path: {model_path}")
+        return local_esm_pretrained.load_model_and_alphabet(str(model_path))
+
+    esm_utils.esm_registry["esm2_3B"] = load_local_esm2_3b
+    simplefold_module = sys.modules.get("models.simplefold.simplefold")
+    if simplefold_module is not None:
+        simplefold_module.esm_registry["esm2_3B"] = load_local_esm2_3b
+
+
+@task_wrapper
+def train(cfg):
+    seed = cfg.get("seed", 42)
+    pl.seed_everything(seed, workers=True)
+    configure_local_esm_registry(cfg.model.get("esm_model", None))
+
+    log.info(f"Instantiating model <{cfg.model._target_}>")
+    model: LightningModule = hydra.utils.instantiate(cfg.model)
+    load_ckpt_path = cfg.get("load_ckpt_path", None)
+
+    if load_ckpt_path is not None:
+        # load existing ckpt
+        log.info(f"Resuming from checkpoint <{cfg.load_ckpt_path}>...")
+        model.strict_loading = False
+
+        # manually reset these variables in case of fine-tuning
+        model.lddt_weight_schedule = cfg.model.get("lddt_weight_schedule", False)
+        model.plddt_training = cfg.model.get("plddt_training", False)
+
+        # reset ESM model to avoid issues in loading FSDP checkpoint
+        model.reset_esm(cfg.model.esm_model)
+
+    log.info(f"Instantiating datamodule <{cfg.data._target_}>")
+    datamodule: LightningDataModule = hydra.utils.instantiate(cfg.data)
+
+    log.info("Instantiating callbacks...")
+    callbacks = instantiate_callbacks(cfg.get("callbacks"))
+
+    log.info("Instantiating loggers...")
+    OmegaConf.set_struct(cfg.logger, True)
+    loggers = instantiate_loggers(cfg.get("logger"))
+
+    log.info(f"Instantiating trainer <{cfg.trainer._target_}>")
+    trainer = instantiate_trainer(
+        cfg.trainer, callbacks=callbacks, logger=loggers, plugins=None
+    )
+
+    object_dict = {
+        "cfg": cfg,
+        "datamodule": datamodule,
+        "model": model,
+        "callbacks": callbacks,
+        "logger": loggers,
+        "trainer": trainer,
+    }
+
+    if log:
+        log.info("Logging hyperparameters!")
+        log_hyperparameters(object_dict)
+
+    log.info("Starting training!")
+    trainer.fit(
+        model=model,
+        datamodule=datamodule,
+        ckpt_path=load_ckpt_path,
+    )
+
+
+@hydra.main(version_base="1.3", config_path="../config", config_name="train.yaml")
+def submit_run(cfg):
+    OmegaConf.resolve(cfg)
+    extras(cfg)
+    create_folders(cfg)
+    train(cfg)
+    return
+
+
+if __name__ == "__main__":
+    submit_run()

scripts/train_fsdp.py

@@ -0,0 +1,161 @@
+#
+# For licensing see accompanying LICENSE file.
+# Copyright (c) 2025 Apple Inc. Licensed under MIT License.
+#
+
+import os
+import hydra
+import torch
+import functools
+from omegaconf import OmegaConf
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+import lightning.pytorch as pl
+from lightning.pytorch import LightningDataModule, LightningModule
+from lightning.pytorch.strategies import FSDPStrategy
+from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
+
+from models.simplefold.torch.blocks import DiTBlock
+from onescience.utils.simplefold.utils import (
+    extras,
+    create_folders,
+    task_wrapper,
+)
+from onescience.utils.simplefold.instantiators import (
+    instantiate_callbacks,
+    instantiate_loggers,
+)
+from onescience.utils.simplefold.logging_utils import log_hyperparameters
+from onescience.utils.simplefold.pylogger import RankedLogger
+
+log = RankedLogger(__name__, rank_zero_only=True)
+ESM2_MODEL_NAME = "esm2_t36_3B_UR50D"
+torch.set_float32_matmul_precision("medium")
+
+
+def configure_local_esm_registry(esm_model: str | None) -> None:
+    if esm_model != "esm2_3B":
+        return
+
+    model_path = Path(
+        os.getenv(
+            "SIMPLEFOLD_ESM2_MODEL_PATH",
+            ROOT / "weight" / "esm_models" / f"{ESM2_MODEL_NAME}.pt",
+        )
+    )
+    regression_path = model_path.with_name(f"{model_path.stem}-contact-regression.pt")
+
+    if not model_path.exists() or not regression_path.exists():
+        raise FileNotFoundError(
+            "Local ESM-2 3B weights were not found. Expected files:\n"
+            f"  {model_path}\n"
+            f"  {regression_path}\n"
+            "Set SIMPLEFOLD_ESM2_MODEL_PATH to the local .pt file if it lives elsewhere."
+        )
+
+    import models.esm.pretrained as local_esm_pretrained
+    import onescience.utils.simplefold.esm_utils as esm_utils
+
+    def load_local_esm2_3b():
+        print(f"Loading ESM-2 3B weights from local path: {model_path}")
+        return local_esm_pretrained.load_model_and_alphabet(str(model_path))
+
+    esm_utils.esm_registry["esm2_3B"] = load_local_esm2_3b
+    simplefold_module = sys.modules.get("models.simplefold.simplefold")
+    if simplefold_module is not None:
+        simplefold_module.esm_registry["esm2_3B"] = load_local_esm2_3b
+
+
+@task_wrapper
+def train(cfg):
+    seed = cfg.get("seed", 42)
+    pl.seed_everything(seed, workers=True)
+    configure_local_esm_registry(cfg.model.get("esm_model", None))
+
+    log.info(f"Instantiating model <{cfg.model._target_}>")
+    model: LightningModule = hydra.utils.instantiate(cfg.model)
+    load_ckpt_path = cfg.get("load_ckpt_path", None)
+
+    if load_ckpt_path is not None:
+        # load existing ckpt
+        log.info(f"Resuming from checkpoint <{cfg.load_ckpt_path}>...")
+        model.strict_loading = False
+
+        # manually reset these variables in case of fine-tuning
+        model.lddt_weight_schedule = cfg.model.get("lddt_weight_schedule", False)
+        model.plddt_training = cfg.model.get("plddt_training", False)
+
+    log.info(f"Instantiating datamodule <{cfg.data._target_}>")
+    datamodule: LightningDataModule = hydra.utils.instantiate(cfg.data)
+
+    log.info("Instantiating callbacks...")
+    callbacks = instantiate_callbacks(cfg.get("callbacks"))
+
+    log.info("Instantiating loggers...")
+    OmegaConf.set_struct(cfg.logger, True)
+    loggers = instantiate_loggers(cfg.get("logger"))
+
+    # When using FSDP, we need to manually specify the wrap policy
+    # and activation checkpointing policy for transformer layers.
+
+    log.info(f"Instantiating trainer <{cfg.trainer._target_}>")
+
+    esm_layer_class = model.esm_model.layers[0].__class__
+
+    transformer_auto_wrapper_policy = functools.partial(
+        transformer_auto_wrap_policy,
+        transformer_layer_cls={DiTBlock, esm_layer_class},
+    )
+    strategy = FSDPStrategy(
+        auto_wrap_policy=transformer_auto_wrapper_policy,
+        activation_checkpointing_policy={DiTBlock, esm_layer_class},
+        use_orig_params=True,
+        state_dict_type="sharded",
+        limit_all_gathers=True,
+        cpu_offload=False
+    )
+    trainer = hydra.utils.instantiate(
+        cfg.trainer, 
+        strategy=strategy,
+        callbacks=callbacks, 
+        logger=loggers, 
+        plugins=None
+    )
+
+    object_dict = {
+        "cfg": cfg,
+        "datamodule": datamodule,
+        "model": model,
+        "callbacks": callbacks,
+        "logger": loggers,
+        "trainer": trainer,
+    }
+
+    if log:
+        log.info("Logging hyperparameters!")
+        log_hyperparameters(object_dict)
+
+    log.info("Starting training!")
+    trainer.fit(
+        model=model,
+        datamodule=datamodule,
+        ckpt_path=load_ckpt_path,
+    )
+
+
+@hydra.main(version_base="1.3", config_path="../config", config_name="base_train.yaml")
+def submit_run(cfg):
+    OmegaConf.resolve(cfg)
+    extras(cfg)
+    create_folders(cfg)
+    train(cfg)
+    return
+
+
+if __name__ == "__main__":
+    submit_run()