Upload folder using huggingface_hub

README.md

@@ -1,122 +1,247 @@
----
-library_name: transformers
-tags:
-  - biology
-  - protein-structure
-  - esmfold2
-  - multimodal-protein-model
----
-
-# FastPLMs ESMFold2
-
-FastPLMs ESMFold2 is a self-contained Hugging Face `AutoModel` wrapper for Biohub's ESMFold2 and ESMFold2-Fast structure predictors. It vendors the released Biohub ESMFold2 model code, ESMC backbone code, input builder, MSA helpers, and structure export utilities needed for remote-code loading.
-
-## Load With AutoModel
+---
+library_name: transformers
+tags:
+  - biology
+  - protein-structure
+  - esmfold2
+  - multimodal-protein-model
+---
+
+# FastPLMs ESMFold2
+
+FastPLMs ESMFold2 is a self-contained Hugging Face `AutoModel` wrapper for
+Biohub's ESMFold2, ESMFold2-Fast, and experimental ESMFold2 structure
+predictors. It vendors the released Biohub ESMFold2 model code, input builder,
+MSA helpers, and structure export utilities, while loading the PLM backbone
+through FastPLMs ESM++.
+
+## Load With AutoModel
+
+```python
+import torch
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained(
+    "Synthyra/ESMFold2-Fast",
+    trust_remote_code=True,
+    dtype=torch.float32,
+).eval().cuda()
+```
+
+Use `Synthyra/ESMFold2` for the full model, `Synthyra/ESMFold2-Fast` for the
+faster release variant, and the `Synthyra/ESMFold2-Experimental*` checkpoints
+for differentiable binder design and experimental critic ensembles.
+The folding trunk runs in fp32; the 6B FastPLMs ESM++ backbone is loaded in
+bf16 by default via `esmc_precision="bf16"` and uses the flex attention backend
+by default inside ESMFold2.
+
+## Fold One Protein
+
+```python
+sequence = "MKTLLILAVVAAALA"
+
+result = model.fold_protein(
+    sequence,
+    num_loops=3,
+    num_sampling_steps=50,
+    num_diffusion_samples=1,
+    seed=0,
+)
+
+print(float(result.plddt.mean()))
+print(float(result.ptm))
+```
+
+## Experimental Test-Time Training
+
+TTT is disabled by default. Standard `fold_protein(...)`, `fold(...)`, raw tensor
+inference, and `state_dict()` keys are unchanged unless you explicitly pass
+`ttt=True` or call `fold_protein_ttt(...)`.
+
+The ESMFold2 TTT path is experimental and protein-only in v1. It trains local
+LoRA adapters only on `_esmc` with a masked language modeling objective. The
+folding trunk, confidence head, diffusion head, and structure input pipeline are
+frozen. TTT can improve difficult low-confidence folds, but it adds substantial
+test-time compute and can degrade already confident predictions.
+
+```python
+result = model.fold_protein(
+    "MSTNPKPQRKTKRNT",
+    num_loops=1,
+    num_sampling_steps=10,
+    num_diffusion_samples=1,
+    seed=0,
+    ttt=True,
+    ttt_config={
+        "steps": 1,
+        "ags": 1,
+        "batch_size": 1,
+        "lora_rank": 8,
+        "lora_alpha": 32.0,
+    },
+)
+
+print(result.ttt_metrics["losses"])
+print(result.ttt_metrics["step_plddts"])
+print(result.ttt_metrics["best_step"])
+```
+
+`load_esmc=True` is required for TTT because the ESM++ MLM head is loaded lazily
+from `config.esmc_id`. If that pretrained MLM head cannot be loaded, TTT raises
+an assertion instead of silently using a random head.
+
+## Save mmCIF or PDB
+
+```python
+model.save_as_cif(result, "prediction.cif")
+model.save_as_pdb(result, "prediction.pdb")
+
+cif_text = model.result_to_cif(result)
+pdb_text = model.result_to_pdb(result)
+```
+
+`result_to_cif` preserves the full `MolecularComplex`. `result_to_pdb` converts through Biohub's protein-only `ProteinComplex` representation, so use mmCIF for complexes with ligands or nucleic acids.
+
+## Fold Complexes
+
+```python
+types = model.input_types
+
+complex_input = types.StructurePredictionInput(
+    sequences=[
+        types.ProteinInput(id="A", sequence="MKTLLILAVVAAALA"),
+        types.DNAInput(id="B", sequence="GATAGC"),
+        types.LigandInput(id="L", ccd=["SAH"]),
+    ]
+)
+
+result = model.fold(
+    complex_input,
+    num_loops=3,
+    num_sampling_steps=50,
+    num_diffusion_samples=1,
+    seed=0,
+)
+
+model.save_as_cif(result, "complex_prediction.cif")
+```
+
+## Binder Design With FastPLMs ESMFold2
+
+FastPLMs includes a FastPLMs-only port of the Biohub ESMFold2 binder design
+tutorial at `cookbook/tutorials/binder_design_fastplms.py`. The workflow uses
+ESMFold2 experimental checkpoints for differentiable folding losses, ESM++ for
+sequence regularization, and ESMFold2 hero critics for final confidence scoring.
+
+![FastPLMs EGFR minibinder design](https://raw.githubusercontent.com/Synthyra/FastPLMs/main/docs/assets/egfr_fastplms_binder_design.png)
+
+The optimizer follows the official strategy:
+
+1. Optimize mutable `#` residues as continuous amino acid logits.
+2. Suppress cysteine design by masking cysteine logits and gradients.
+3. Backpropagate through ESMFold2 `res_type_soft` using intra-contact,
+   inter-contact, and globularity losses from the distogram.
+4. Add an ESM++ masked-LM pseudoperplexity regularizer on mutable binder
+   residues.
+5. Keep the late-trajectory sequence with the best iPTM.
+6. Fold the selected sequence with the final critic ensemble and write
+   `results.parquet`, `selection.parquet`, `trajectory.jsonl`,
+   `best_sequences.fasta`, and per-critic PDB/CIF/logit files.
+
+Run the verified EGFR 128 amino acid de novo minibinder example:
+
+```bash
+cd /home/ubuntu/FastPLMs
+
+sudo -n docker run --gpus all --rm \
+  -v /home/ubuntu/FastPLMs:/app \
+  -v /home/ubuntu/FastPLMs:/workspace \
+  -v /home/ubuntu/.cache/huggingface:/workspace/.cache/huggingface \
+  -w /workspace fastplms-esmfold2 \
+  python /app/cookbook/tutorials/binder_design_fastplms.py \
+    --backend local \
+    --target-name egfr \
+    --binder-sequence '################################################################################################################################' \
+    --not-antibody \
+    --steps 150 \
+    --batch-size 1 \
+    --seed 103 \
+    --output-dir /workspace/campaign_egfr_len128_b1_s150_seed103_consensus_cli
+```
+
+Verified result:
+
+| Metric | Value |
+| :--- | :--- |
+| Binder length | `128` |
+| Seed | `103` |
+| Steps | `150` |
+| Hero mean iPTM | `0.913870` |
+| Hero min iPTM | `0.904600` |
+| All four hero critics above 0.9 | `True` |
+
+Binder sequence:
+
+```text
+SAVKHLLEIVKYLEEAIEKALEVDPVFLVPPAAEELLIAAKVIKELAKENPELIEVYELLMKAVKGLKKLVRSNDKEILREVIRLLRKAAKVIREILKNNPDLDPELRKALEELAKVLEEIAEVLEQQ
+```
+
+See the full guide in [`docs/binder_design.md`](https://github.com/Synthyra/FastPLMs/blob/main/docs/binder_design.md)
+for Modal execution, official pI and selection scoring, per-critic metrics, and
+the tested cheaper step-count boundary.
+
+## Use MSAs
+
+```python
+types = model.input_types
+
+msa = types.MSA.from_a3m("query.a3m", max_sequences=128)
+input_with_msa = types.StructurePredictionInput(
+    sequences=[
+        types.ProteinInput(id="A", sequence=msa.query, msa=msa),
+    ]
+)
+
+result = model.fold(input_with_msa, num_sampling_steps=50, seed=0)
+```
+
+## Raw Tensor Inference
+
+```python
+features, chain_infos = model.prepare_structure_input(complex_input, seed=0)
+
+with torch.inference_mode():
+    output = model(
+        **features,
+        num_loops=3,
+        num_sampling_steps=50,
+        num_diffusion_samples=1,
+    )
+
+decoded = model.input_builder.decode(output, features, chain_infos)
+```
+
+Set `load_esmc=False` when loading if you want to provide precomputed `lm_hidden_states` manually or run folding-trunk tests without loading the 6B ESM++ backbone:
 
 ```python
-import torch
-from transformers import AutoModel
-
 model = AutoModel.from_pretrained(
     "Synthyra/ESMFold2-Fast",
-    trust_remote_code=True,
-    dtype=torch.float32,
-).eval().cuda()
-```
-
-Use `Synthyra/ESMFold2` for the full model and `Synthyra/ESMFold2-Fast` for the faster release variant.
-The folding trunk runs in fp32; the 6B ESMC backbone is loaded in bf16 by default via `esmc_precision="bf16"`.
-
-## Fold One Protein
-
-```python
-sequence = "MKTLLILAVVAAALA"
-
-result = model.fold_protein(
-    sequence,
-    num_loops=3,
-    num_sampling_steps=50,
-    num_diffusion_samples=1,
-    seed=0,
-)
-
-print(float(result.plddt.mean()))
-print(float(result.ptm))
-```
-
-## Save mmCIF or PDB
-
-```python
-model.save_as_cif(result, "prediction.cif")
-model.save_as_pdb(result, "prediction.pdb")
-
-cif_text = model.result_to_cif(result)
-pdb_text = model.result_to_pdb(result)
-```
-
-`result_to_cif` preserves the full `MolecularComplex`. `result_to_pdb` converts through Biohub's protein-only `ProteinComplex` representation, so use mmCIF for complexes with ligands or nucleic acids.
-
-## Fold Complexes
-
-```python
-types = model.input_types
-
-complex_input = types.StructurePredictionInput(
-    sequences=[
-        types.ProteinInput(id="A", sequence="MKTLLILAVVAAALA"),
-        types.DNAInput(id="B", sequence="GATAGC"),
-        types.LigandInput(id="L", ccd=["SAH"]),
-    ]
-)
-
-result = model.fold(
-    complex_input,
-    num_loops=3,
-    num_sampling_steps=50,
-    num_diffusion_samples=1,
-    seed=0,
-)
-
-model.save_as_cif(result, "complex_prediction.cif")
-```
-
-## Use MSAs
-
-```python
-types = model.input_types
-
-msa = types.MSA.from_a3m("query.a3m", max_sequences=128)
-input_with_msa = types.StructurePredictionInput(
-    sequences=[
-        types.ProteinInput(id="A", sequence=msa.query, msa=msa),
-    ]
-)
-
-result = model.fold(input_with_msa, num_sampling_steps=50, seed=0)
-```
-
-## Raw Tensor Inference
-
-```python
-features, chain_infos = model.prepare_structure_input(complex_input, seed=0)
-
-with torch.inference_mode():
-    output = model(
-        **features,
-        num_loops=3,
-        num_sampling_steps=50,
-        num_diffusion_samples=1,
-    )
-
-decoded = model.input_builder.decode(output, features, chain_infos)
+    trust_remote_code=True,
+    load_esmc=False,
+).cuda().eval()
 ```
 
-Set `load_esmc=False` when loading if you want to provide precomputed `lm_hidden_states` manually or run folding-trunk tests without loading the 6B ESMC backbone:
+For FP8 LM inference, install `transformer_engine.pytorch` in a CUDA
+environment with FP8-capable hardware and load the shared FastPLMs ESM++
+backbone with:
 
 ```python
 model = AutoModel.from_pretrained(
     "Synthyra/ESMFold2-Fast",
     trust_remote_code=True,
-    load_esmc=False,
+    esmc_precision="fp8",
 ).cuda().eval()
 ```
+
+FP8 is inference-only for the ESMFold2 LM backbone. TTT remains a bf16/fp32
+path.

__init__.py

@@ -1,4 +1,5 @@
-from .configuration_esmfold2 import ESMFold2Config
-from .modeling_esmfold2 import ESMFold2Model
-
-__all__ = ["ESMFold2Config", "ESMFold2Model"]
+from .configuration_esmfold2 import ESMFold2Config
+from .modeling_esmfold2_experimental import ESMFold2ExperimentalModel
+from .modeling_esmfold2 import ESMFold2Model
+
+__all__ = ["ESMFold2Config", "ESMFold2ExperimentalModel", "ESMFold2Model"]

configuration_esmfold2.py

@@ -1,194 +1,206 @@
-# Copyright 2026 Biohub. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""ESMFold2 model configuration."""
-
-from __future__ import annotations
-
-from dataclasses import asdict, dataclass, field
-
-from transformers.configuration_utils import PretrainedConfig
-
-# ---------------------------------------------------------------------------
-# Nested dataclass configs
-# ---------------------------------------------------------------------------
-
-_DEFAULT_ESMC_HF_REPO = "biohub/ESMC-6B"
-
-
-@dataclass
-class MSAEncoderConfig:
-    """Config for the optional MSA encoder module (Large MSA models only)."""
-
-    enabled: bool = False
-    d_msa: int = 128
-    d_hidden: int = 32
-    n_layers: int = 4
-    n_heads_msa: int = 8
-    msa_head_width: int = 32
-
-
-@dataclass
-class ParcaeConfig:
-    """Release-only config for the parcae diffusion-loop scheduler."""
-
-    enabled: bool = True
-    poisson_mean: float = 3.0
-    min_steps: int = 1
-    max_steps: int | None = 6
-    coda_n_layers: int = 2
-
-
-@dataclass
-class LMEncoderConfig:
-    """Release-only config for the LM-side pair encoder."""
-
-    enabled: bool = True
-    n_layers: int = 4
-    lm_dropout: float = 0.25
-    per_loop_lm_dropout: bool = True
-
-
-@dataclass
-class AtomAttentionConfig:
-    """Config for SWA atom encoder/decoder with 3D RoPE."""
-
-    d_atom: int = 128
-    d_token: int = 768
-    n_blocks: int = 3
-    n_heads: int = 4
-    swa_window_size: int = 128
-    expansion_ratio: int = 2
-    # 3D RoPE config
-    spatial_rope_base_frequency: float = 20.0
-    n_spatial_rope_pairs_per_axis: int = 2
-    n_uid_rope_pairs: int = 10
-    uid_rope_base_frequency: float = 10000.0
-
-
-@dataclass
-class FoldingTrunkConfig:
-    n_layers: int = 24
-    n_heads: int = 8
-    dropout: float = 0.0
-
-
-@dataclass
-class InputsEmbedderConfig:
-    d_inputs: int = 451
-    atom_encoder: AtomAttentionConfig = field(default_factory=AtomAttentionConfig)
-
-    def __post_init__(self):
-        if isinstance(self.atom_encoder, dict):
-            self.atom_encoder = AtomAttentionConfig(**self.atom_encoder)
-
-
-@dataclass
-class DiffusionModuleConfig:
-    """Config for the DiffusionModule."""
-
-    sigma_data: float = 16.0
-    c_atom: int = 128
-    c_token: int = 768
-    c_z: int = 256
-    c_s_inputs: int = 451
-    fourier_dim: int = 256
-    relpos_r_max: int = 32
-    relpos_s_max: int = 2
-    atom_num_blocks: int = 3
-    atom_num_heads: int = 4
-    token_num_blocks: int = 12
-    token_num_heads: int = 16
-    transition_multiplier: int = 2
-
-
-@dataclass
-class DiffusionStructureHeadConfig:
-    """Config for the diffusion-based structure prediction head."""
-
-    diffusion_module: DiffusionModuleConfig = field(
-        default_factory=DiffusionModuleConfig
-    )
-    distogram_bins: int = 128
-
-    # Training noise: sigma ~ sigma_data * exp(mu + sigma * N(0,1))
-    train_noise_log_mean: float = -1.2
-    train_noise_log_std: float = 1.5
-
-    # Sampling defaults (ODE)
-    gamma_0: float = 0.605
-    gamma_min: float = 1.107
-    noise_scale: float = 0.0
-    step_scale: float = 1.0
-
-    # Inference schedule defaults
-    inference_s_max: float = 160.0
-    inference_s_min: float = 4e-4
-    inference_p: float = 8.0
-    inference_num_steps: int = 68
-
-    def __post_init__(self):
-        if isinstance(self.diffusion_module, dict):
-            self.diffusion_module = DiffusionModuleConfig(**self.diffusion_module)
-
-
-@dataclass
-class ConfidenceHeadConfig:
-    enabled: bool = True
-    num_plddt_bins: int = 50
-    num_pde_bins: int = 64
-    num_pae_bins: int = 64
-    min_dist: float = 2.0
-    max_dist: float = 52.0
-    distogram_bins: int = 128
-    folding_trunk: FoldingTrunkConfig = field(
-        default_factory=lambda: FoldingTrunkConfig(n_layers=4)
-    )
-
-    def __post_init__(self):
-        if isinstance(self.folding_trunk, dict):
-            self.folding_trunk = FoldingTrunkConfig(**self.folding_trunk)
-
-
-# ---------------------------------------------------------------------------
-# Top-level config
-# ---------------------------------------------------------------------------
-
-
-class ESMFold2Config(PretrainedConfig):
-    """
-    Configuration for the ESMFold2 structure prediction model.
-
-    Uses SWA atom encoders with 3D RoPE, a diffusion transformer,
-    a folding trunk, and an ESMC 6B PLM backbone.
-
-    Configuration objects inherit from [`PretrainedConfig`] and can be used to control
-    the model outputs. Read the documentation from [`PretrainedConfig`] for more
-    information.
-
-    Args:
-        d_single (`int`, defaults to 384):
-            Dimensionality of single (per-residue) representations.
-        d_pair (`int`, defaults to 256):
-            Dimensionality of pair (residue-residue) representations.
-        n_relative_residx_bins (`int`, defaults to 32):
-            Number of bins for relative residue index encoding.
-        n_relative_chain_bins (`int`, defaults to 2):
-            Number of bins for relative chain encoding.
-        num_loops (`int`, defaults to 10):
-            Number of trunk loops for iterative refinement.
-        num_diffusion_samples (`int`, defaults to 8):
-            Number of parallel structure predictions to generate.
+# Copyright 2026 Biohub. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""ESMFold2 model configuration."""
+
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass, field
+
+from transformers.configuration_utils import PretrainedConfig
+
+# ---------------------------------------------------------------------------
+# Nested dataclass configs
+# ---------------------------------------------------------------------------
+
+_DEFAULT_ESMC_HF_REPO = "Synthyra/ESMplusplus_6B"
+_DEFAULT_ESMC_ATTN_BACKEND = "flex"
+_ESMC_ID_ALIASES = {
+    "biohub/ESMC-300M": "Synthyra/ESMplusplus_small",
+    "biohub/ESMC-600M": "Synthyra/ESMplusplus_large",
+    "biohub/ESMC-6B": "Synthyra/ESMplusplus_6B",
+}
+
+
+def normalize_esmc_id(esmc_id: str) -> str:
+    if esmc_id in _ESMC_ID_ALIASES:
+        return _ESMC_ID_ALIASES[esmc_id]
+    return esmc_id
+
+
+@dataclass
+class MSAEncoderConfig:
+    """Config for the optional MSA encoder module (Large MSA models only)."""
+
+    enabled: bool = False
+    d_msa: int = 128
+    d_hidden: int = 32
+    n_layers: int = 4
+    n_heads_msa: int = 8
+    msa_head_width: int = 32
+
+
+@dataclass
+class ParcaeConfig:
+    """Release-only config for the parcae diffusion-loop scheduler."""
+
+    enabled: bool = True
+    poisson_mean: float = 3.0
+    min_steps: int = 1
+    max_steps: int | None = 6
+    coda_n_layers: int = 2
+
+
+@dataclass
+class LMEncoderConfig:
+    """Release-only config for the LM-side pair encoder."""
+
+    enabled: bool = True
+    n_layers: int = 4
+    lm_dropout: float = 0.25
+    per_loop_lm_dropout: bool = True
+
+
+@dataclass
+class AtomAttentionConfig:
+    """Config for SWA atom encoder/decoder with 3D RoPE."""
+
+    d_atom: int = 128
+    d_token: int = 768
+    n_blocks: int = 3
+    n_heads: int = 4
+    swa_window_size: int = 128
+    expansion_ratio: int = 2
+    # 3D RoPE config
+    spatial_rope_base_frequency: float = 20.0
+    n_spatial_rope_pairs_per_axis: int = 2
+    n_uid_rope_pairs: int = 10
+    uid_rope_base_frequency: float = 10000.0
+
+
+@dataclass
+class FoldingTrunkConfig:
+    n_layers: int = 24
+    n_heads: int = 8
+    dropout: float = 0.0
+
+
+@dataclass
+class InputsEmbedderConfig:
+    d_inputs: int = 451
+    atom_encoder: AtomAttentionConfig = field(default_factory=AtomAttentionConfig)
+
+    def __post_init__(self):
+        if isinstance(self.atom_encoder, dict):
+            self.atom_encoder = AtomAttentionConfig(**self.atom_encoder)
+
+
+@dataclass
+class DiffusionModuleConfig:
+    """Config for the DiffusionModule."""
+
+    sigma_data: float = 16.0
+    c_atom: int = 128
+    c_token: int = 768
+    c_z: int = 256
+    c_s_inputs: int = 451
+    fourier_dim: int = 256
+    relpos_r_max: int = 32
+    relpos_s_max: int = 2
+    atom_num_blocks: int = 3
+    atom_num_heads: int = 4
+    token_num_blocks: int = 12
+    token_num_heads: int = 16
+    transition_multiplier: int = 2
+
+
+@dataclass
+class DiffusionStructureHeadConfig:
+    """Config for the diffusion-based structure prediction head."""
+
+    diffusion_module: DiffusionModuleConfig = field(
+        default_factory=DiffusionModuleConfig
+    )
+    distogram_bins: int = 128
+
+    # Training noise: sigma ~ sigma_data * exp(mu + sigma * N(0,1))
+    train_noise_log_mean: float = -1.2
+    train_noise_log_std: float = 1.5
+
+    # Sampling defaults (ODE)
+    gamma_0: float = 0.605
+    gamma_min: float = 1.107
+    noise_scale: float = 0.0
+    step_scale: float = 1.0
+
+    # Inference schedule defaults
+    inference_s_max: float = 160.0
+    inference_s_min: float = 4e-4
+    inference_p: float = 8.0
+    inference_num_steps: int = 68
+
+    def __post_init__(self):
+        if isinstance(self.diffusion_module, dict):
+            self.diffusion_module = DiffusionModuleConfig(**self.diffusion_module)
+
+
+@dataclass
+class ConfidenceHeadConfig:
+    enabled: bool = True
+    num_plddt_bins: int = 50
+    num_pde_bins: int = 64
+    num_pae_bins: int = 64
+    min_dist: float = 2.0
+    max_dist: float = 52.0
+    distogram_bins: int = 128
+    folding_trunk: FoldingTrunkConfig = field(
+        default_factory=lambda: FoldingTrunkConfig(n_layers=4)
+    )
+
+    def __post_init__(self):
+        if isinstance(self.folding_trunk, dict):
+            self.folding_trunk = FoldingTrunkConfig(**self.folding_trunk)
+
+
+# ---------------------------------------------------------------------------
+# Top-level config
+# ---------------------------------------------------------------------------
+
+
+class ESMFold2Config(PretrainedConfig):
+    """
+    Configuration for the ESMFold2 structure prediction model.
+
+    Uses SWA atom encoders with 3D RoPE, a diffusion transformer,
+    a folding trunk, and an ESMC 6B PLM backbone.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control
+    the model outputs. Read the documentation from [`PretrainedConfig`] for more
+    information.
+
+    Args:
+        d_single (`int`, defaults to 384):
+            Dimensionality of single (per-residue) representations.
+        d_pair (`int`, defaults to 256):
+            Dimensionality of pair (residue-residue) representations.
+        n_relative_residx_bins (`int`, defaults to 32):
+            Number of bins for relative residue index encoding.
+        n_relative_chain_bins (`int`, defaults to 2):
+            Number of bins for relative chain encoding.
+        num_loops (`int`, defaults to 10):
+            Number of trunk loops for iterative refinement.
+        num_diffusion_samples (`int`, defaults to 8):
+            Number of parallel structure predictions to generate.
         lm_dropout (`float`, defaults to 0.0):
             Dropout probability on LM pair embeddings. When > 0, dropout is
             applied with ``training=True`` (including at inference) to match
@@ -196,103 +208,121 @@ class ESMFold2Config(PretrainedConfig):
         force_lm_dropout_during_inference (`bool`, defaults to False):
             When True, apply ``lm_dropout`` even when ``model.eval()`` and
             ``lm_dropout`` > 0. Binder-design loads set this to True.
+        lm_mask_pct (`float`, defaults to 0.0):
+            Fraction of LM residue tokens randomly replaced with the LM mask
+            token before running the PLM backbone.
         disable_msa_features (`bool`, defaults to False):
             When True, zero out MSA-derived ``profile`` and ``deletion_mean``
             before the inputs embedder (experimental medium/large checkpoints).
-        inputs (`InputsEmbedderConfig`):
-            Configuration for the inputs embedder module.
-        folding_trunk (`FoldingTrunkConfig`):
-            Configuration for the folding trunk.
-        structure_head (`DiffusionStructureHeadConfig`):
-            Configuration for the diffusion-based structure prediction head.
-        confidence_head (`ConfidenceHeadConfig`):
-            Configuration for the confidence prediction head.
-
-    Examples:
-
-    ```python
-    >>> from transformers import ESMFold2Config, ESMFold2ExperimentalModel
-
-    >>> # Initializing an ESMFold2 configuration
-    >>> configuration = ESMFold2Config(type="experimental")
-
-    >>> # Initializing a model (with random weights) from the configuration
-    >>> model = ESMFold2ExperimentalModel(configuration)
-
-    >>> # Accessing the model configuration
-    >>> configuration = model.config
-    ```
-    """
-
-    model_type = "esmfold2"
-    has_no_defaults_at_init = True
-
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-        self.type: str = kwargs.get("type", "release")
-        if self.type not in ("release", "experimental"):
-            raise ValueError(
-                f"ESMFold2Config.type must be 'release' or 'experimental', "
-                f"got {self.type!r}"
-            )
-
-        # Top-level scalar fields
-        self.d_single: int = kwargs.get("d_single", 384)
-        self.d_pair: int = kwargs.get("d_pair", 256)
-        self.n_relative_residx_bins: int = kwargs.get("n_relative_residx_bins", 32)
-        self.n_relative_chain_bins: int = kwargs.get("n_relative_chain_bins", 2)
-        self.num_loops: int = kwargs.get("num_loops", 10)
-        self.num_diffusion_samples: int = kwargs.get("num_diffusion_samples", 8)
-        # If True, ``profile`` / ``deletion_mean`` are zeroed before the inputs
-        # embedder.
-        self.disable_msa_features: bool = kwargs.get("disable_msa_features", False)
-        self.lm_dropout: float = kwargs.get("lm_dropout", 0.0)
+        inputs (`InputsEmbedderConfig`):
+            Configuration for the inputs embedder module.
+        folding_trunk (`FoldingTrunkConfig`):
+            Configuration for the folding trunk.
+        structure_head (`DiffusionStructureHeadConfig`):
+            Configuration for the diffusion-based structure prediction head.
+        confidence_head (`ConfidenceHeadConfig`):
+            Configuration for the confidence prediction head.
+
+    Examples:
+
+    ```python
+    >>> from transformers import ESMFold2Config, ESMFold2ExperimentalModel
+
+    >>> # Initializing an ESMFold2 configuration
+    >>> configuration = ESMFold2Config(type="experimental")
+
+    >>> # Initializing a model (with random weights) from the configuration
+    >>> model = ESMFold2ExperimentalModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```
+    """
+
+    model_type = "esmfold2"
+    has_no_defaults_at_init = True
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+        self.type: str = kwargs.get("type", "release")
+        if self.type not in ("release", "experimental"):
+            raise ValueError(
+                f"ESMFold2Config.type must be 'release' or 'experimental', "
+                f"got {self.type!r}"
+            )
+
+        # Top-level scalar fields
+        self.d_single: int = kwargs.get("d_single", 384)
+        self.d_pair: int = kwargs.get("d_pair", 256)
+        self.n_relative_residx_bins: int = kwargs.get("n_relative_residx_bins", 32)
+        self.n_relative_chain_bins: int = kwargs.get("n_relative_chain_bins", 2)
+        self.num_loops: int = kwargs.get("num_loops", 10)
+        self.num_diffusion_samples: int = kwargs.get("num_diffusion_samples", 8)
+        # If True, ``profile`` / ``deletion_mean`` are zeroed before the inputs
+        # embedder.
+        self.disable_msa_features: bool = kwargs.get("disable_msa_features", False)
+        self.lm_dropout: float = kwargs.get("lm_dropout", 0.0)
         self.force_lm_dropout_during_inference: bool = kwargs.get(
             "force_lm_dropout_during_inference", False
         )
+        self.lm_mask_pct: float = kwargs.get("lm_mask_pct", 0.0)
 
         self.lm_d_model: int = kwargs.get("lm_d_model", 2560)
-        self.lm_num_layers: int = kwargs.get("lm_num_layers", 80)
-        # Required, no default — every shipped HF export must name its ESMC backbone.
-        self.esmc_id: str = kwargs.get("esmc_id", _DEFAULT_ESMC_HF_REPO)
-
-        def _init_nested(cls, val):
-            if isinstance(val, cls):
-                return val
-            if isinstance(val, dict):
-                return cls(**val)
-            return cls()
-
-        self.inputs = _init_nested(InputsEmbedderConfig, kwargs.get("inputs"))
-        self.folding_trunk = _init_nested(
-            FoldingTrunkConfig, kwargs.get("folding_trunk")
-        )
-        self.structure_head = _init_nested(
-            DiffusionStructureHeadConfig, kwargs.get("structure_head")
-        )
-        self.confidence_head = _init_nested(
-            ConfidenceHeadConfig, kwargs.get("confidence_head")
-        )
-        self.msa_encoder = _init_nested(MSAEncoderConfig, kwargs.get("msa_encoder"))
-        # Release-only modules — ignored when ``type == "experimental"``.
-        self.parcae = _init_nested(ParcaeConfig, kwargs.get("parcae"))
-        self.lm_encoder = _init_nested(LMEncoderConfig, kwargs.get("lm_encoder"))
-        # If True, MSA encoder output replaces the pair stream; if False, it is added.
-        self.msa_encoder_overwrite: bool = bool(
-            kwargs.get("msa_encoder_overwrite", True)
-        )
-
-    def to_dict(self):
-        output = super().to_dict()
-        output["inputs"] = asdict(self.inputs)
-        output["folding_trunk"] = asdict(self.folding_trunk)
-        output["structure_head"] = asdict(self.structure_head)
-        output["confidence_head"] = asdict(self.confidence_head)
-        output["msa_encoder"] = asdict(self.msa_encoder)
-        output["parcae"] = asdict(self.parcae)
-        output["lm_encoder"] = asdict(self.lm_encoder)
-        return output
-
-
-__all__ = ["ESMFold2Config", "MSAEncoderConfig", "ParcaeConfig", "LMEncoderConfig"]
+        self.lm_num_layers: int = kwargs.get("lm_num_layers", 80)
+        # Backward-compatible field name; values now point to FastPLMs ESM++.
+        raw_esmc_id = (
+            kwargs["esmc_id"] if "esmc_id" in kwargs else _DEFAULT_ESMC_HF_REPO
+        )
+        self.esmc_id: str = normalize_esmc_id(raw_esmc_id)
+        self.esmc_attn_backend: str = (
+            kwargs["esmc_attn_backend"]
+            if "esmc_attn_backend" in kwargs
+            else _DEFAULT_ESMC_ATTN_BACKEND
+        )
+
+        def _init_nested(cls, val):
+            if isinstance(val, cls):
+                return val
+            if isinstance(val, dict):
+                return cls(**val)
+            return cls()
+
+        self.inputs = _init_nested(InputsEmbedderConfig, kwargs.get("inputs"))
+        self.folding_trunk = _init_nested(
+            FoldingTrunkConfig, kwargs.get("folding_trunk")
+        )
+        self.structure_head = _init_nested(
+            DiffusionStructureHeadConfig, kwargs.get("structure_head")
+        )
+        self.confidence_head = _init_nested(
+            ConfidenceHeadConfig, kwargs.get("confidence_head")
+        )
+        self.msa_encoder = _init_nested(MSAEncoderConfig, kwargs.get("msa_encoder"))
+        # Release-only modules — ignored when ``type == "experimental"``.
+        self.parcae = _init_nested(ParcaeConfig, kwargs.get("parcae"))
+        self.lm_encoder = _init_nested(LMEncoderConfig, kwargs.get("lm_encoder"))
+        # If True, MSA encoder output replaces the pair stream; if False, it is added.
+        self.msa_encoder_overwrite: bool = bool(
+            kwargs.get("msa_encoder_overwrite", True)
+        )
+
+    def to_dict(self):
+        output = super().to_dict()
+        output["inputs"] = asdict(self.inputs)
+        output["folding_trunk"] = asdict(self.folding_trunk)
+        output["structure_head"] = asdict(self.structure_head)
+        output["confidence_head"] = asdict(self.confidence_head)
+        output["msa_encoder"] = asdict(self.msa_encoder)
+        output["parcae"] = asdict(self.parcae)
+        output["lm_encoder"] = asdict(self.lm_encoder)
+        return output
+
+
+__all__ = [
+    "ESMFold2Config",
+    "MSAEncoderConfig",
+    "ParcaeConfig",
+    "LMEncoderConfig",
+    "normalize_esmc_id",
+]

esmfold2_affine3d.py

@@ -558,4 +558,3 @@ def build_affine3d_from_coordinates(
     )
 
     return affine, coord_mask
-

esmfold2_aligner.py

@@ -99,4 +99,3 @@ class Aligner:
         aligned_atom37_positions[mobile.atom37_mask] = aligned_atom_tensor
 
         return replace(mobile, atom37_positions=aligned_atom37_positions)
-

esmfold2_atom_indexer.py

@@ -13,4 +13,3 @@ class AtomIndexer:
         return index_by_atom_name(
             getattr(self.structure, self.property), atom_names, self.dim
         )
-

esmfold2_conformers.py

@@ -289,4 +289,3 @@ def get_ccd_leaving_atoms(comp_id: str) -> set[str]:
 
     _CCD_LEAVING_ATOMS_CACHE[comp_id] = leaving_atoms
     return leaving_atoms
-

esmfold2_constants.py

@@ -560,4 +560,3 @@ RNA_BACKBONE_ATOMS = [
     "O2'",
     "C1'",
 ]
-

esmfold2_constants_esm3.py

@@ -135,4 +135,3 @@ KEYWORDS_IDF = IN_REPO_DATA_FOLDER / "keyword_idf_safety_filtered_58641.npy"
 
 RESID_CSV = "data/uniref90_and_mgnify90_residue_annotations_gt_1k_proteins.csv"
 INTERPRO2KEYWORDS = IN_REPO_DATA_FOLDER / "interpro_29026_to_keywords_58641.csv"
-

esmfold2_input_builder.py

@@ -252,4 +252,3 @@ def deserialize_structure_prediction_input(
         distogram_conditioning=distogram_conditioning,
         covalent_bonds=covalent_bonds,
     )
-

esmfold2_metrics.py

@@ -5,7 +5,7 @@ from einops import rearrange
 from torch import Tensor
 from torch.amp import autocast  # type: ignore
 
-from . import esmfold2_residue_constants as residue_constants
+from . import esmfold2_residue_constants as residue_constants
 from .esmfold2_misc import binpack, unbinpack
 from .esmfold2_protein_structure import (
     compute_alignment_tensors,
@@ -371,4 +371,3 @@ def compute_gdt_ts(
     return compute_gdt_ts_no_alignment(
         rotated_mobile, centered_target, atom_exists_mask, reduction
     )
-

esmfold2_misc.py

@@ -502,4 +502,3 @@ def concat_objects(objs: Sequence[Any], separator: Any | None = None):
                 return torch.cat(objs)  # type: ignore
         case _:
             raise TypeError(type(objs[0]))
-

esmfold2_mmcif_parsing.py

@@ -10,7 +10,7 @@ from typing import Union
 import biotite.structure as bs
 import biotite.structure.io.pdbx as pdbx
 
-from . import esmfold2_residue_constants as residue_constants
+from . import esmfold2_residue_constants as residue_constants
 
 # Define PathOrBuffer for the opensource version
 PathOrBuffer = Union[str, os.PathLike, io.StringIO]
@@ -467,4 +467,3 @@ class MmcifWrapper:
             return self._parse_nonpoly_from_mmcif()
         except Exception:
             return self._parse_nonpoly_fallback()
-

esmfold2_molecular_complex.py

@@ -23,7 +23,7 @@ from biotite.structure.io.pdbx import (
     set_structure,
 )
 
-from . import esmfold2_residue_constants as residue_constants
+from . import esmfold2_residue_constants as residue_constants
 from .esmfold2_metrics import compute_lddt, compute_rmsd
 from .esmfold2_protein_complex import ProteinComplex, ProteinComplexMetadata
 
@@ -44,6 +44,7 @@ class MolecularComplexResult:
     residue_index: torch.Tensor | None = None
     entity_id: torch.Tensor | None = None
     sae_features: np.ndarray | None = None  # [L, n_features]
+    ttt_metrics: dict[str, Any] | None = None
 
 
 @dataclass
@@ -1223,4 +1224,3 @@ class MolecularComplex:
         return cls.from_state_dict(
             msgpack.loads(brotli.decompress(bytes), strict_map_key=False)
         )
-

esmfold2_msa.py

@@ -504,4 +504,3 @@ class FastMSA(SequentialDataclass):
             if headers is not None:
                 all_headers.extend(headers)
         return cls(np.concatenate(arrays, axis=0), all_headers)
-

esmfold2_msa_filter_sequences.py

@@ -80,4 +80,3 @@ def hhfilter(
         with output_file.open() as f:
             indices = [int(line[1:].strip()) for line in f if line.startswith(">")]
         return indices
-

esmfold2_normalize_coordinates.py

@@ -77,4 +77,3 @@ def apply_frame_to_coords(coords: Tensor, frame: Affine3D) -> Tensor:
 
 def normalize_coordinates(coords: Tensor) -> Tensor:
     return apply_frame_to_coords(coords, get_protein_normalization_frame(coords))
-

esmfold2_output.py

@@ -222,4 +222,3 @@ def build_molecular_complex(
             assembly_composition=None,
         ),
     )
-

esmfold2_paired_msa.py

@@ -243,4 +243,3 @@ def construct_paired_msa(
         paired_mask[:, token_mask] = chain_paired[:, None]
 
     return msa_residues, deletion_value, paired_mask
-

esmfold2_parsing.py

@@ -110,4 +110,3 @@ def write_sequences(sequences: Iterable[tuple[str, str]], path: PathOrBuffer) ->
     finally:
         if needs_closing:
             handle.close()  # type: ignore
-

esmfold2_predicted_aligned_error.py

@@ -102,4 +102,3 @@ def tm_loss(
         loss = loss.mean()
 
     return loss
-

esmfold2_prepare_input.py

@@ -1461,4 +1461,3 @@ def prepare_esmfold2_input(
     chains, tokens, atoms = build_chains_from_input(input, seed)
     features = build_feature_tensors(chains, tokens, atoms, input)
     return features, chains
-

esmfold2_processor.py

@@ -353,4 +353,3 @@ class ESMFold2InputBuilder:
 
 
 __all__ = ["ESMFold2InputBuilder", "clean_esmfold2_input"]
-

esmfold2_protein_chain.py

@@ -20,7 +20,7 @@ from biotite.structure.io.pdbx import set_structure as set_structure_pdbx
 from scipy.spatial import ConvexHull, KDTree
 from scipy.spatial.distance import cdist, pdist, squareform
 
-from . import esmfold2_residue_constants as residue_constants
+from . import esmfold2_residue_constants as residue_constants
 from .esmfold2_misc import slice_python_object_as_numpy
 from .esmfold2_affine3d import Affine3D
 from .esmfold2_aligner import Aligner
@@ -1373,4 +1373,3 @@ class ProteinChain:
         )  # (1, L)
 
         return coordinates, plddt, residue_index
-

esmfold2_protein_complex.py

@@ -27,7 +27,7 @@ from biotite.structure.io.pdbx.convert import _get_transformations, get_structur
 from biotite.structure.util import matrix_rotate
 from scipy.spatial import KDTree
 
-from . import esmfold2_residue_constants as residue_constants
+from . import esmfold2_residue_constants as residue_constants
 from .esmfold2_misc import slice_python_object_as_numpy
 from .esmfold2_affine3d import Affine3D
 from .esmfold2_aligner import Aligner
@@ -1238,4 +1238,3 @@ def protein_chain_to_protein_complex(chain: ProteinChain) -> ProteinComplex:
         for i, chain in enumerate(complex_chains)
     ]
     return ProteinComplex.from_chains(complex_chains)
-

esmfold2_protein_structure.py

@@ -8,7 +8,7 @@ import torch.nn.functional as F
 from torch import Tensor
 from torch.amp import autocast  # type: ignore
 
-from . import esmfold2_residue_constants as residue_constants
+from . import esmfold2_residue_constants as residue_constants
 from .esmfold2_misc import unbinpack
 from .esmfold2_affine3d import Affine3D
 
@@ -304,4 +304,3 @@ def compute_gdt_ts_no_alignment(
         return score
     else:
         raise ValueError("Unrecognized reduction: '{reduction}'")
-

esmfold2_residue_constants.py

@@ -1221,4 +1221,3 @@ def _make_restype_atom14_to_atom37():
 RESTYPE_ATOM14_TO_ATOM37 = _make_restype_atom14_to_atom37()
 RESTYPE_ATOM37_TO_ATOM14 = _make_restype_atom37_to_atom14()
 CHAIN_BREAK_TOKEN = "|"
-

esmfold2_sequential_dataclass.py

@@ -155,4 +155,3 @@ class SequentialDataclass(ABC):
             items[0],  # type: ignore
             **updated_fields,
         )
-

esmfold2_system.py

@@ -43,4 +43,3 @@ def run_subprocess_with_errorcheck(
             f"Command failed with errorcode {e.returncode}." f"\n\n{e.stderr.decode()}"
         )
     return p
-

esmfold2_types.py

@@ -31,4 +31,3 @@ __all__ = [
     "CovalentBond",
     "StructurePredictionInput",
 ]
-

esmfold2_utils_types.py

@@ -31,4 +31,3 @@ class FunctionAnnotation:
     def __len__(self) -> int:
         """Length of the annotation."""
         return self.end - self.start + 1
-

modeling_esmfold2_experimental.py

@@ -0,0 +1,1001 @@
+"""FastPLMs ESMFold2 experimental architecture.
+
+This module supports Biohub's experimental binder-design checkpoints. The
+released ESMFold2 architecture in ``modeling_esmfold2.py`` intentionally
+rejects those configs because the experimental trunk uses explicit pair-loop
+re-injection and a different confidence/MSA stack.
+"""
+
+from __future__ import annotations
+
+import math
+from pathlib import Path
+from typing import Any, cast
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from transformers.modeling_utils import PreTrainedModel
+
+from .configuration_esmfold2 import ESMFold2Config
+from .modeling_esmfold2 import (
+    _load_fastplms_esmplusplus_for_esmfold2,
+    _lm_precision_context,
+)
+from .modeling_esmfold2_common import (
+    CHAR_VOCAB_SIZE,
+    MAX_ATOMIC_NUMBER,
+    NUM_RES_TYPES,
+    DiffusionModule,
+    DiffusionStructureHead,
+    DiffusionTransformer,
+    FoldingTrunk,
+    InputsEmbedder,
+    LanguageModelShim,
+    MSAPairWeightedAveraging,
+    OuterProductMean,
+    PairUpdateBlock,
+    ResIdxAsymIdSymIdEntityIdEncoding,
+    RowAttentionPooling,
+    SwiGLUMLP,
+    TriangleMultiplicativeUpdate,
+    _categorical_mean,
+    _compute_intra_token_idx,
+    _seed_context,
+    compute_lm_hidden_states,
+    gather_rep_atom_coords,
+    gather_token_to_atom,
+)
+
+_EPS = 1e-5
+_NONPOLYMER_ID = 3
+
+
+class ConfidenceHead(nn.Module):
+    """Experimental confidence head predicting pLDDT, PAE, pTM, and ipTM."""
+
+    boundaries: Tensor
+
+    def __init__(self, config: ESMFold2Config) -> None:
+        super().__init__()
+        ch = config.confidence_head
+        d_single = config.d_single
+        d_pair = config.d_pair
+        d_inputs = config.inputs.d_inputs
+
+        boundaries = torch.linspace(ch.min_dist, ch.max_dist, ch.distogram_bins - 1)
+        self.register_buffer("boundaries", boundaries)
+        self.dist_bin_pairwise_embed = nn.Embedding(ch.distogram_bins, d_pair)
+
+        self.s_norm = nn.LayerNorm(d_single)
+        self.s_inputs_to_single = nn.Linear(d_inputs, d_single, bias=False)
+        self.s_to_z = nn.Linear(d_inputs, d_pair, bias=False)
+        self.s_to_z_transpose = nn.Linear(d_inputs, d_pair, bias=False)
+        self.s_to_z_prod_in1 = nn.Linear(d_inputs, d_pair, bias=False)
+        self.s_to_z_prod_in2 = nn.Linear(d_inputs, d_pair, bias=False)
+        self.s_to_z_prod_out = nn.Linear(d_pair, d_pair, bias=False)
+        self.s_input_to_s = nn.Linear(d_inputs, d_single, bias=False)
+        self.s_inputs_norm = nn.LayerNorm(d_inputs)
+        self.z_norm = nn.LayerNorm(d_pair)
+        self.row_attention_pooling = RowAttentionPooling(
+            d_pair=d_pair, d_single=d_single
+        )
+
+        pf = ch.folding_trunk
+        self.folding_trunk = FoldingTrunk(
+            n_layers=pf.n_layers, d_pair=d_pair, expansion_ratio=4
+        )
+
+        self.plddt_ln = nn.LayerNorm(d_single)
+        max_atoms_per_token = 23
+        self.plddt_weight = nn.Parameter(
+            torch.zeros(max_atoms_per_token, d_single, ch.num_plddt_bins)
+        )
+        self.pae_head = nn.Linear(d_pair, ch.num_pae_bins, bias=False)
+
+    def set_kernel_backend(self, backend: str | None) -> None:
+        self.folding_trunk.set_kernel_backend(backend)
+
+    def set_chunk_size(self, chunk_size: int | None) -> None:
+        self.folding_trunk.set_chunk_size(chunk_size)
+
+    @staticmethod
+    def _repeat_batch(x: Tensor, num_diffusion_samples: int) -> Tensor:
+        if num_diffusion_samples == 1:
+            return x
+        return x.repeat_interleave(num_diffusion_samples, 0)
+
+    @staticmethod
+    def _flatten_sample_axis(x: Tensor) -> Tensor:
+        if x.ndim == 4:
+            b, mult, n, c = x.shape
+            return x.reshape(b * mult, n, c)
+        return x
+
+    def forward(
+        self,
+        s_inputs: Tensor,
+        z: Tensor,
+        x_pred: Tensor,
+        distogram_atom_idx: Tensor,
+        token_attention_mask: Tensor,
+        atom_to_token: Tensor,
+        atom_attention_mask: Tensor,
+        asym_id: Tensor,
+        mol_type: Tensor,
+        num_diffusion_samples: int = 1,
+        relative_position_encoding: Tensor | None = None,
+        token_bonds_encoding: Tensor | None = None,
+    ) -> dict[str, Tensor]:
+        s_inputs_normed = self.s_inputs_norm(s_inputs)
+        z_base = self.z_norm(z)
+        if relative_position_encoding is not None:
+            z_base = z_base + relative_position_encoding
+        if token_bonds_encoding is not None:
+            z_base = z_base + token_bonds_encoding
+        z_base = z_base + self.s_to_z(s_inputs_normed).unsqueeze(2)
+        z_base = z_base + self.s_to_z_transpose(s_inputs_normed).unsqueeze(1)
+        z_base = z_base + self.s_to_z_prod_out(
+            self.s_to_z_prod_in1(s_inputs_normed)[:, :, None, :]
+            * self.s_to_z_prod_in2(s_inputs_normed)[:, None, :, :]
+        )
+
+        pair = self._repeat_batch(z_base, num_diffusion_samples)
+        x_pred_flat = self._flatten_sample_axis(x_pred)
+        atom_to_token_m = self._repeat_batch(atom_to_token, num_diffusion_samples)
+        atom_mask_m = self._repeat_batch(atom_attention_mask, num_diffusion_samples)
+        rep_idx_m = self._repeat_batch(distogram_atom_idx, num_diffusion_samples).long()
+        mask = self._repeat_batch(token_attention_mask, num_diffusion_samples)
+        batch_mult = pair.shape[0]
+
+        rep_coords = gather_rep_atom_coords(x_pred_flat, rep_idx_m)
+        rep_distances = torch.cdist(
+            rep_coords, rep_coords, compute_mode="donot_use_mm_for_euclid_dist"
+        )
+        distogram_bins = (
+            (rep_distances.unsqueeze(-1) > self.boundaries).sum(dim=-1).long()
+        )
+        pair = pair + self.dist_bin_pairwise_embed(distogram_bins)
+
+        pair_mask = mask[:, :, None].float() * mask[:, None, :].float()
+        pair = pair + self.folding_trunk(pair, pair_attention_mask=pair_mask)
+        single = self.row_attention_pooling(pair, mask)
+
+        atom_mask_f = atom_mask_m.float()
+        s_at_atoms = gather_token_to_atom(single, atom_to_token_m)
+        s_at_atoms = self.plddt_ln(s_at_atoms)
+        intra_idx = _compute_intra_token_idx(atom_to_token_m)
+        intra_idx = intra_idx.clamp(max=self.plddt_weight.shape[0] - 1)
+        plddt_weight = self.plddt_weight[intra_idx]
+        plddt_logits = torch.einsum("...c,...cb->...b", s_at_atoms, plddt_weight)
+        plddt_per_atom = _categorical_mean(plddt_logits, start=0.0, end=1.0)
+
+        length = single.shape[1]
+        plddt_sum = torch.zeros(
+            batch_mult, length, device=single.device, dtype=plddt_per_atom.dtype
+        )
+        atom_count = torch.zeros(
+            batch_mult, length, device=single.device, dtype=plddt_per_atom.dtype
+        )
+        atom_mask_t = atom_mask_f.to(plddt_per_atom.dtype)
+        plddt_sum.scatter_add_(1, atom_to_token_m, plddt_per_atom * atom_mask_t)
+        atom_count.scatter_add_(1, atom_to_token_m, atom_mask_t)
+        plddt = plddt_sum / atom_count.clamp(min=1e-6)
+
+        complex_plddt = (plddt_per_atom * atom_mask_f).sum(dim=-1) / (
+            atom_mask_f.sum(dim=-1) + _EPS
+        )
+
+        expanded_type = self._repeat_batch(mol_type, num_diffusion_samples)
+        expanded_asym = self._repeat_batch(asym_id, num_diffusion_samples)
+        is_ligand = (expanded_type == _NONPOLYMER_ID).float()
+        inter_chain = (
+            expanded_asym.unsqueeze(-1) != expanded_asym.unsqueeze(-2)
+        ).float()
+        near_contact = (rep_distances < 8).float()
+        interface_per_token = (
+            near_contact * inter_chain * (1.0 - is_ligand).unsqueeze(-1)
+        ).amax(dim=-1)
+        iplddt_weight = torch.where(
+            is_ligand.bool(),
+            torch.full_like(interface_per_token, 2.0),
+            interface_per_token,
+        )
+        iplddt_weight_atoms = gather_token_to_atom(
+            iplddt_weight.unsqueeze(-1), atom_to_token_m
+        ).squeeze(-1)
+        atom_iplddt_w = atom_mask_f * iplddt_weight_atoms
+        complex_iplddt = (plddt_per_atom * atom_iplddt_w).sum(dim=-1) / (
+            atom_iplddt_w.sum(dim=-1) + _EPS
+        )
+        plddt_ca = plddt_per_atom.gather(1, rep_idx_m)
+
+        pae_logits = self.pae_head(pair)
+        pae = _categorical_mean(pae_logits, start=0.0, end=32.0).detach()
+
+        n_bins = pae_logits.shape[-1]
+        bin_width = 32.0 / n_bins
+        bin_centers = torch.arange(
+            0.5 * bin_width, 32.0, bin_width, device=pae_logits.device
+        )
+        mask_f = mask.float()
+        n_res = mask_f.sum(dim=-1, keepdim=True)
+        d0 = 1.24 * (n_res.clamp(min=19) - 15) ** (1 / 3) - 1.8
+        tm_per_bin = 1 / (1 + (bin_centers / d0) ** 2)
+        pae_probs = F.softmax(pae_logits, dim=-1)
+        tm_expected = (pae_probs * tm_per_bin[:, None, None, :]).sum(dim=-1)
+
+        pair_mask_2d = mask_f.unsqueeze(-1) * mask_f.unsqueeze(-2)
+        ptm_per_row = (tm_expected * pair_mask_2d).sum(dim=-1) / (
+            pair_mask_2d.sum(dim=-1) + _EPS
+        )
+        ptm = ptm_per_row.max(dim=-1).values
+
+        inter_chain_mask = (
+            expanded_asym.unsqueeze(-1) != expanded_asym.unsqueeze(-2)
+        ).float() * pair_mask_2d
+        iptm_per_row = (tm_expected * inter_chain_mask).sum(dim=-1) / (
+            inter_chain_mask.sum(dim=-1) + _EPS
+        )
+        iptm = iptm_per_row.max(dim=-1).values
+
+        max_chain_id = int(expanded_asym.max().item()) if batch_mult > 0 else 0
+        n_chains = max_chain_id + 1
+        pair_chains_iptm = torch.zeros(
+            batch_mult,
+            n_chains,
+            n_chains,
+            device=tm_expected.device,
+            dtype=tm_expected.dtype,
+        )
+        for c1 in range(n_chains):
+            chain_c1 = (expanded_asym == c1).float() * mask_f
+            if chain_c1.sum() == 0:
+                continue
+            for c2 in range(n_chains):
+                chain_c2 = (expanded_asym == c2).float() * mask_f
+                pair_m = chain_c1.unsqueeze(-1) * chain_c2.unsqueeze(-2)
+                denom = pair_m.sum(dim=(-1, -2)) + _EPS
+                pair_chains_iptm[:, c1, c2] = (tm_expected * pair_m).sum(
+                    dim=(-1, -2)
+                ) / denom
+
+        return {
+            "plddt_logits": plddt_logits,
+            "plddt": plddt.detach(),
+            "plddt_per_atom": plddt_per_atom.detach(),
+            "plddt_ca": plddt_ca.detach(),
+            "complex_plddt": complex_plddt.detach(),
+            "complex_iplddt": complex_iplddt.detach(),
+            "pae_logits": pae_logits,
+            "pae": pae,
+            "ptm": ptm.detach(),
+            "iptm": iptm.detach(),
+            "pair_chains_iptm": pair_chains_iptm.detach(),
+        }
+
+
+class _TransitionFFN(nn.Module):
+    def __init__(self, d_model: int, expansion_ratio: int = 4) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm(d_model)
+        self.ffn = SwiGLUMLP(d_model, expansion_ratio=expansion_ratio, bias=False)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.ffn(self.norm(x))
+
+
+class MSAEncoderBlock(nn.Module):
+    """One experimental MSA update block."""
+
+    def __init__(
+        self,
+        d_msa: int,
+        d_pair: int,
+        d_hidden: int = 32,
+        n_heads_msa: int = 8,
+        msa_head_width: int = 32,
+    ) -> None:
+        super().__init__()
+        self.outer_product_mean = OuterProductMean(
+            d_msa, d_hidden, d_pair, divide_outer_before_proj=True
+        )
+        self.msa_pair_weighted_averaging = MSAPairWeightedAveraging(
+            d_msa, d_pair, n_heads_msa, msa_head_width
+        )
+        self.msa_transition = _TransitionFFN(d_msa, expansion_ratio=4)
+        self.tri_mul_out = TriangleMultiplicativeUpdate(dim=d_pair, _outgoing=True)
+        self.tri_mul_in = TriangleMultiplicativeUpdate(dim=d_pair, _outgoing=False)
+        self.pair_transition = _TransitionFFN(d_pair, expansion_ratio=4)
+
+    def set_chunk_size(self, chunk_size: int | None) -> None:
+        self.outer_product_mean.set_chunk_size(chunk_size)
+        self.tri_mul_out.set_chunk_size(chunk_size)
+        self.tri_mul_in.set_chunk_size(chunk_size)
+
+    def forward(
+        self,
+        msa_repr: Tensor,
+        pair_repr: Tensor,
+        msa_attention_mask: Tensor,
+        pair_attention_mask: Tensor,
+        msa_track_mask: Tensor | None = None,
+    ) -> tuple[Tensor, Tensor]:
+        mask4d = (
+            msa_track_mask[:, None, None, None].to(dtype=msa_repr.dtype)
+            if msa_track_mask is not None
+            else None
+        )
+
+        pair_mask4d = mask4d[:, :, :1] if mask4d is not None else None
+
+        msa_update = self.msa_pair_weighted_averaging(
+            msa_repr, pair_repr, pair_attention_mask
+        )
+        if mask4d is not None:
+            msa_update = msa_update * mask4d
+        msa_repr = msa_repr + msa_update
+
+        msa_transition = self.msa_transition(msa_repr)
+        if mask4d is not None:
+            msa_transition = msa_transition * mask4d
+        msa_repr = msa_repr + msa_transition
+
+        pair_opm = self.outer_product_mean(msa_repr, msa_attention_mask)
+        if pair_mask4d is not None:
+            pair_opm = pair_opm * pair_mask4d
+        pair_repr = pair_repr + pair_opm
+
+        pair_out = self.tri_mul_out(pair_repr, mask=pair_attention_mask)
+        if pair_mask4d is not None:
+            pair_out = pair_out * pair_mask4d
+        pair_repr = pair_repr + pair_out
+
+        pair_in = self.tri_mul_in(pair_repr, mask=pair_attention_mask)
+        if pair_mask4d is not None:
+            pair_in = pair_in * pair_mask4d
+        pair_repr = pair_repr + pair_in
+
+        pair_transition = self.pair_transition(pair_repr)
+        if pair_mask4d is not None:
+            pair_transition = pair_transition * pair_mask4d
+        pair_repr = pair_repr + pair_transition
+        return msa_repr, pair_repr
+
+
+class MSAEncoder(nn.Module):
+    def __init__(
+        self,
+        d_msa: int,
+        d_pair: int,
+        d_inputs: int,
+        d_hidden: int = 32,
+        n_layers: int = 4,
+        n_heads_msa: int = 8,
+        msa_head_width: int = 32,
+    ) -> None:
+        super().__init__()
+        self.embed = nn.Linear(35, d_msa, bias=False)
+        self.project_inputs = nn.Linear(d_inputs, d_msa, bias=False)
+        self.blocks = nn.ModuleList(
+            [
+                MSAEncoderBlock(
+                    d_msa=d_msa,
+                    d_pair=d_pair,
+                    d_hidden=d_hidden,
+                    n_heads_msa=n_heads_msa,
+                    msa_head_width=msa_head_width,
+                )
+                for _ in range(n_layers)
+            ]
+        )
+
+    def set_chunk_size(self, chunk_size: int | None) -> None:
+        for block in self.blocks:
+            cast(MSAEncoderBlock, block).set_chunk_size(chunk_size)
+
+    def forward(
+        self,
+        x_pair: Tensor,
+        x_inputs: Tensor,
+        msa_oh: Tensor,
+        has_deletion: Tensor,
+        deletion_value: Tensor,
+        msa_attention_mask: Tensor,
+    ) -> Tensor:
+        batch_size, _, depth = msa_attention_mask.shape
+        m_feat = torch.cat(
+            [msa_oh, has_deletion.unsqueeze(-1), deletion_value.unsqueeze(-1)],
+            dim=-1,
+        )
+        m = self.embed(m_feat) + self.project_inputs(x_inputs).unsqueeze(2)
+        if depth > 1:
+            msa_track_mask = msa_attention_mask[:, :, 1:].any(dim=(1, 2))
+        else:
+            msa_track_mask = torch.zeros(
+                batch_size, dtype=torch.bool, device=x_pair.device
+            )
+        tok_mask = msa_attention_mask[:, :, 0]
+        pair_attention_mask = tok_mask.unsqueeze(2) * tok_mask.unsqueeze(1)
+        for block in self.blocks:
+            m, x_pair = cast(MSAEncoderBlock, block)(
+                m,
+                x_pair,
+                msa_attention_mask,
+                pair_attention_mask,
+                msa_track_mask,
+            )
+        return x_pair * msa_track_mask[:, None, None, None].to(dtype=x_pair.dtype)
+
+
+class ESMFold2ExperimentalModel(PreTrainedModel):
+    """Experimental ESMFold2 architecture used by binder-design checkpoints."""
+
+    config_class = ESMFold2Config
+    _keys_to_ignore_on_load_unexpected = [r"\._extra_state$"]
+
+    def __init__(self, config: ESMFold2Config) -> None:
+        super().__init__(config)
+        d_inputs = config.inputs.d_inputs
+        d_pair = config.d_pair
+
+        self.inputs_embedder = InputsEmbedder(config)
+        self.z_init_1 = nn.Linear(d_inputs, d_pair, bias=False)
+        self.z_init_2 = nn.Linear(d_inputs, d_pair, bias=False)
+        self.rel_pos = ResIdxAsymIdSymIdEntityIdEncoding(
+            n_relative_residx_bins=config.n_relative_residx_bins,
+            n_relative_chain_bins=config.n_relative_chain_bins,
+            d_pair=d_pair,
+        )
+        self.token_bonds = nn.Linear(1, d_pair, bias=False)
+        self.language_model = LanguageModelShim(
+            d_z=d_pair, d_model=config.lm_d_model, num_layers=config.lm_num_layers
+        )
+        self._esmc: nn.Module | None = None
+        self._esmc_fp8 = False
+        self._esmfold2_input_builder: Any | None = None
+
+        pf = config.folding_trunk
+        self.folding_trunk = FoldingTrunk(
+            n_layers=pf.n_layers, d_pair=d_pair, expansion_ratio=4
+        )
+        self.pair_loop_proj = nn.Sequential(
+            nn.LayerNorm(d_pair), nn.Linear(d_pair, d_pair, bias=False)
+        )
+        nn.init.zeros_(cast(nn.Linear, self.pair_loop_proj[1]).weight)
+
+        self.structure_head = DiffusionStructureHead(config)
+        self.distogram_head = nn.Linear(
+            d_pair, config.structure_head.distogram_bins, bias=True
+        )
+        self.confidence_head: ConfidenceHead | None = (
+            ConfidenceHead(config) if config.confidence_head.enabled else None
+        )
+
+        msa_cfg = config.msa_encoder
+        self.msa_encoder: MSAEncoder | None = None
+        if msa_cfg.enabled:
+            self.msa_encoder = MSAEncoder(
+                d_msa=msa_cfg.d_msa,
+                d_pair=d_pair,
+                d_inputs=d_inputs,
+                d_hidden=msa_cfg.d_hidden,
+                n_layers=msa_cfg.n_layers,
+                n_heads_msa=msa_cfg.n_heads_msa,
+                msa_head_width=msa_cfg.msa_head_width,
+            )
+
+        self.post_init()
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.parameters()).device
+
+    def set_kernel_backend(self, backend: str | None) -> None:
+        self.folding_trunk.set_kernel_backend(backend)
+        if self.confidence_head is not None:
+            self.confidence_head.set_kernel_backend(backend)
+        self.structure_head.set_kernel_backend(backend)
+
+    def set_chunk_size(self, chunk_size: int | None) -> None:
+        self.folding_trunk.set_chunk_size(chunk_size)
+        if self.confidence_head is not None:
+            self.confidence_head.set_chunk_size(chunk_size)
+        if self.msa_encoder is not None:
+            self.msa_encoder.set_chunk_size(chunk_size)
+
+    def configure_lm_dropout(
+        self,
+        lm_dropout: float,
+        *,
+        force_lm_dropout_during_inference: bool = True,
+    ) -> None:
+        self.config.lm_dropout = lm_dropout
+        self.config.force_lm_dropout_during_inference = (
+            force_lm_dropout_during_inference
+        )
+
+    def load_esmc(self, esmc_model_path: str, precision: str = "bf16") -> None:
+        dtype_map = {
+            "bf16": torch.bfloat16,
+            "fp32": torch.float32,
+        }
+        if precision not in dtype_map:
+            if precision == "fp8":
+                raise RuntimeError(
+                    "esmc_precision='fp8' is supported only by the standard "
+                    "released ESMFold2 model. The experimental binder-design "
+                    "model keeps the FastPLMs ESM++ backbone in bf16 or fp32."
+                )
+            raise ValueError(f"precision must be one of {list(dtype_map)}, got {precision!r}")
+        esmc = _load_fastplms_esmplusplus_for_esmfold2(
+            esmc_model_path=esmc_model_path,
+            attn_backend=self.config.esmc_attn_backend,
+            device=self.device,
+            dtype=dtype_map[precision],
+        )
+        assert esmc.config.hidden_size == self.config.lm_d_model, (
+            f"ESMFold2 expected lm_d_model={self.config.lm_d_model}, "
+            f"but loaded ESM++ hidden_size={esmc.config.hidden_size}."
+        )
+        assert esmc.config.num_hidden_layers == self.config.lm_num_layers, (
+            f"ESMFold2 expected lm_num_layers={self.config.lm_num_layers}, "
+            f"but loaded ESM++ num_hidden_layers={esmc.config.num_hidden_layers}."
+        )
+        for parameter in esmc.parameters():
+            parameter.requires_grad_(False)
+        self._esmc_fp8 = False
+        self._esmc = esmc
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path,
+        *model_args,
+        load_esmc: bool = True,
+        **kwargs,
+    ):
+        if "config" not in kwargs:
+            kwargs["config"] = ESMFold2Config.from_pretrained(
+                pretrained_model_name_or_path, **kwargs
+            )
+        esmc_precision = kwargs.pop("esmc_precision", "bf16")
+        model = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+        if load_esmc:
+            model.load_esmc(model.config.esmc_id, precision=esmc_precision)
+        return model
+
+    def apply_torch_compile(
+        self, mode: str = "fixed_seqlen", dynamic: bool | None = None
+    ) -> None:
+        import torch._dynamo
+
+        torch._dynamo.config.cache_size_limit = 512
+        torch._dynamo.config.accumulated_cache_size_limit = 512
+        torch._dynamo.config.capture_scalar_outputs = True
+
+        if dynamic is None:
+            dynamic = mode == "dynamic_seqlen"
+        compile_kwargs: dict[str, bool] = {"dynamic": dynamic}
+        compile_targets = (
+            PairUpdateBlock,
+            DiffusionTransformer,
+            DiffusionModule,
+            MSAEncoderBlock,
+        )
+
+        def _maybe_compile(module: nn.Module) -> None:
+            if isinstance(module, compile_targets):
+                module.forward = torch.compile(module.forward, **compile_kwargs)
+
+        self.apply(_maybe_compile)
+
+    def _compute_lm_hidden_states(
+        self,
+        input_ids: Tensor,
+        asym_id: Tensor,
+        residue_index: Tensor,
+        mol_type: Tensor,
+        tok_mask: Tensor,
+    ) -> Tensor:
+        assert self._esmc is not None
+        pad_to = 8 if self._esmc_fp8 else None
+        with _lm_precision_context(self._esmc_fp8):
+            return compute_lm_hidden_states(
+                self._esmc,
+                input_ids,
+                asym_id,
+                residue_index,
+                mol_type,
+                tok_mask,
+                pad_to_multiple=pad_to,
+            )
+
+    def forward(
+        self,
+        token_index: Tensor,
+        residue_index: Tensor,
+        asym_id: Tensor,
+        sym_id: Tensor,
+        entity_id: Tensor,
+        mol_type: Tensor,
+        res_type: Tensor,
+        token_bonds: Tensor,
+        token_attention_mask: Tensor,
+        ref_pos: Tensor,
+        ref_element: Tensor,
+        ref_charge: Tensor,
+        ref_atom_name_chars: Tensor,
+        ref_space_uid: Tensor,
+        atom_attention_mask: Tensor,
+        atom_to_token: Tensor,
+        distogram_atom_idx: Tensor,
+        deletion_mean: Tensor | None = None,
+        msa: Tensor | None = None,
+        has_deletion: Tensor | None = None,
+        deletion_value: Tensor | None = None,
+        msa_attention_mask: Tensor | None = None,
+        input_ids: Tensor | None = None,
+        lm_hidden_states: Tensor | None = None,
+        res_type_soft: Tensor | None = None,
+        num_loops: int | None = None,
+        num_diffusion_samples: int | None = None,
+        num_sampling_steps: int | None = None,
+        early_exit: bool = False,
+        seed: int | None = None,
+        calculate_confidence: bool = True,
+        provide_soft_sequence_to_msa_and_profile: bool = True,
+        noise_scale: float | None = None,
+        step_scale: float | None = None,
+        max_inference_sigma: int | None = None,
+    ) -> dict[str, Tensor]:
+        del noise_scale, step_scale, max_inference_sigma
+        tok_mask = token_attention_mask
+        atm_mask = atom_attention_mask
+        n_loops = num_loops if num_loops is not None else self.config.num_loops
+        n_samples = (
+            num_diffusion_samples
+            if num_diffusion_samples is not None
+            else self.config.num_diffusion_samples
+        )
+
+        if res_type.dim() == 2:
+            res_type_oh = F.one_hot(res_type.long(), num_classes=NUM_RES_TYPES).float()
+            res_type_oh = res_type_oh * tok_mask.unsqueeze(-1).float()
+        else:
+            res_type_oh = res_type.float()
+
+        if msa is not None:
+            msa_oh_profile = F.one_hot(msa.long(), num_classes=NUM_RES_TYPES).float()
+            if msa_attention_mask is not None:
+                mask_f = msa_attention_mask.float().unsqueeze(-1)
+                msa_oh_profile = msa_oh_profile * mask_f
+                valid_seq_count = msa_attention_mask.float().sum(dim=1).clamp(min=1)
+                profile = msa_oh_profile.sum(dim=1) / valid_seq_count.unsqueeze(-1)
+            else:
+                profile = msa_oh_profile.mean(dim=1)
+        else:
+            profile = res_type_oh
+
+        if res_type_soft is not None:
+            res_type_oh = res_type_soft.float()
+            if (
+                not self.config.disable_msa_features
+                and provide_soft_sequence_to_msa_and_profile
+            ):
+                profile = res_type_oh
+                msa = res_type_oh.unsqueeze(1)
+                msa_attention_mask = tok_mask.unsqueeze(1)
+
+        if deletion_mean is None:
+            deletion_mean = torch.zeros(
+                res_type.shape[0], res_type.shape[1], device=res_type.device
+            )
+        if self.config.disable_msa_features:
+            profile = torch.zeros_like(profile)
+            deletion_mean = torch.zeros_like(deletion_mean)
+
+        ref_element_oh = F.one_hot(
+            ref_element.long(), num_classes=MAX_ATOMIC_NUMBER
+        ).float()
+        ref_atom_name_chars_oh = F.one_hot(
+            ref_atom_name_chars.long(), num_classes=CHAR_VOCAB_SIZE
+        ).float()
+        atm_mask_f = atm_mask.float()
+        ref_element_oh = ref_element_oh * atm_mask_f.unsqueeze(-1)
+        ref_atom_name_chars_oh = ref_atom_name_chars_oh * atm_mask_f.unsqueeze(
+            -1
+        ).unsqueeze(-1)
+        atom_to_token = atom_to_token * atm_mask.long()
+
+        use_amp = ref_pos.device.type == "cuda"
+        with (
+            torch.set_grad_enabled(res_type_soft is not None),
+            torch.amp.autocast("cuda", enabled=use_amp, dtype=torch.bfloat16),
+        ):
+            x_inputs = self.inputs_embedder(
+                aatype=res_type_oh,
+                profile=profile.float(),
+                deletion_mean=deletion_mean.float(),
+                ref_pos=ref_pos,
+                atom_attention_mask=atm_mask,
+                ref_space_uid=ref_space_uid,
+                ref_charge=ref_charge,
+                ref_element=ref_element_oh,
+                ref_atom_name_chars=ref_atom_name_chars_oh,
+                atom_to_token=atom_to_token,
+            )
+
+            z_init = self.z_init_1(x_inputs).unsqueeze(2) + self.z_init_2(
+                x_inputs
+            ).unsqueeze(1)
+            relative_position_encoding = self.rel_pos(
+                residue_index=residue_index,
+                asym_id=asym_id,
+                sym_id=sym_id,
+                entity_id=entity_id,
+                token_index=token_index,
+            )
+            token_bonds_encoding = self.token_bonds(token_bonds.float())
+            z_init = z_init + relative_position_encoding + token_bonds_encoding
+
+            if (
+                lm_hidden_states is None
+                and input_ids is not None
+                and self._esmc is not None
+            ):
+                lm_hidden_states = self._compute_lm_hidden_states(
+                    input_ids, asym_id, residue_index, mol_type, tok_mask
+                )
+            if lm_hidden_states is not None:
+                lm_dropout = (
+                    self.config.lm_dropout
+                    if self.config.force_lm_dropout_during_inference or self.training
+                    else 0.0
+                )
+                lm_z = self.language_model(
+                    lm_hidden_states.detach(), lm_dropout=lm_dropout
+                )
+                z_init = z_init + lm_z.to(z_init.dtype)
+
+            msa_kwargs: dict[str, Tensor] | None = None
+            if self.msa_encoder is not None and msa is not None:
+                if msa.dim() == 4:
+                    batch_msa, depth, length_msa, _ = msa.shape
+                    msa_oh = msa.permute(0, 2, 1, 3).float()
+                else:
+                    batch_msa, depth, length_msa = msa.shape
+                    msa_oh = F.one_hot(
+                        msa.permute(0, 2, 1).long(), num_classes=NUM_RES_TYPES
+                    ).float()
+                msa_attn = (
+                    msa_attention_mask.permute(0, 2, 1).float()
+                    if msa_attention_mask is not None
+                    else tok_mask[:, :, None].expand(-1, -1, depth).float()
+                )
+                msa_oh = msa_oh * msa_attn.unsqueeze(-1)
+                hd = (
+                    has_deletion.permute(0, 2, 1).float()
+                    if has_deletion is not None
+                    else torch.zeros(batch_msa, length_msa, depth, device=msa.device)
+                )
+                dv = (
+                    deletion_value.permute(0, 2, 1).float()
+                    if deletion_value is not None
+                    else torch.zeros(batch_msa, length_msa, depth, device=msa.device)
+                )
+                msa_kwargs = {
+                    "x_inputs": x_inputs,
+                    "msa_oh": msa_oh,
+                    "has_deletion": hd,
+                    "deletion_value": dv,
+                    "msa_attention_mask": msa_attn,
+                }
+
+            pair_mask = tok_mask[:, :, None].float() * tok_mask[:, None, :].float()
+            z = torch.zeros_like(z_init)
+            prev_pair: Tensor | None = None
+            prev_disto_probs: Tensor | None = None
+            for loop_num in range(n_loops + 1):
+                z = z_init + self.pair_loop_proj(z)
+                if msa_kwargs is not None and self.msa_encoder is not None:
+                    z = z + self.msa_encoder(x_pair=z, **msa_kwargs).to(z.dtype)
+                z = self.folding_trunk(z, pair_attention_mask=pair_mask)
+
+                if early_exit and loop_num < n_loops:
+                    l2_converged = False
+                    if prev_pair is not None and loop_num > 0:
+                        rel_l2 = (z.float() - prev_pair.float()).norm() / prev_pair.float().norm().clamp(
+                            min=1e-8
+                        )
+                        l2_converged = rel_l2.item() < 0.25
+                    prev_pair = z.detach().clone()
+                    sym_z = z.float() + z.float().transpose(-2, -3)
+                    cur_probs = F.softmax(self.distogram_head(sym_z).float(), dim=-1)
+                    if prev_disto_probs is not None and loop_num > 0:
+                        kl_per_pair = (
+                            cur_probs
+                            * (
+                                cur_probs.clamp(min=1e-8)
+                                / prev_disto_probs.clamp(min=1e-8)
+                            ).log()
+                        ).sum(-1)
+                        kl = (kl_per_pair + kl_per_pair.transpose(-1, -2)).mean() / 2
+                        if l2_converged or kl.item() < 0.05:
+                            break
+                    prev_disto_probs = cur_probs.detach()
+
+            distogram_logits = self.distogram_head(z + z.transpose(-2, -3))
+
+        with torch.no_grad(), _seed_context(seed):
+            structure_output = self.structure_head.sample(
+                z_trunk=z.float(),
+                s_inputs=x_inputs,
+                s_trunk=None,
+                relative_position_encoding=relative_position_encoding,
+                ref_pos=ref_pos,
+                ref_charge=ref_charge,
+                ref_mask=atm_mask,
+                ref_element=ref_element_oh,
+                ref_atom_name_chars=ref_atom_name_chars_oh,
+                ref_space_uid=ref_space_uid,
+                tok_idx=atom_to_token,
+                asym_id=asym_id,
+                residue_index=residue_index,
+                entity_id=entity_id,
+                token_index=token_index,
+                sym_id=sym_id,
+                token_attention_mask=tok_mask,
+                num_diffusion_samples=n_samples,
+                num_sampling_steps=num_sampling_steps,
+                return_atom_repr=False,
+                denoising_early_exit_rmsd=(0.10 if early_exit else None),
+            )
+        sample_coords = structure_output["sample_atom_coords"]
+        assert sample_coords is not None
+
+        output: dict[str, Tensor] = {
+            "distogram_logits": distogram_logits,
+            "sample_atom_coords": sample_coords,
+        }
+        if calculate_confidence and self.confidence_head is not None:
+            confidence_output = self.confidence_head(
+                s_inputs=x_inputs.detach(),
+                z=z.detach().float(),
+                x_pred=sample_coords.detach(),
+                distogram_atom_idx=distogram_atom_idx,
+                token_attention_mask=tok_mask,
+                atom_to_token=atom_to_token,
+                atom_attention_mask=atm_mask,
+                asym_id=asym_id,
+                mol_type=mol_type,
+                num_diffusion_samples=n_samples,
+                relative_position_encoding=relative_position_encoding.detach(),
+                token_bonds_encoding=token_bonds_encoding.detach(),
+            )
+            output.update(confidence_output)
+        output["atom_pad_mask"] = (
+            atm_mask.unsqueeze(0) if atm_mask.dim() == 1 else atm_mask
+        )
+        output["residue_index"] = residue_index
+        output["entity_id"] = entity_id
+        return output
+
+    @property
+    def input_builder(self):
+        if self._esmfold2_input_builder is None:
+            from .esmfold2_processor import ESMFold2InputBuilder
+
+            self._esmfold2_input_builder = ESMFold2InputBuilder()
+        return self._esmfold2_input_builder
+
+    @property
+    def input_types(self):
+        from . import esmfold2_types
+
+        return esmfold2_types
+
+    def prepare_structure_input(self, input, seed: int | None = None):
+        return self.input_builder.prepare_input(input, seed=seed, device=self.device)
+
+    @torch.no_grad()
+    def infer_protein(self, seq: str, **forward_kwargs) -> dict[str, Tensor]:
+        from .protein_utils import prepare_protein_features
+
+        features = prepare_protein_features(seq)
+        features = {name: tensor.to(self.device) for name, tensor in features.items()}
+        output = self(**features, **forward_kwargs)
+        for name in (
+            "res_type",
+            "atom_to_token",
+            "ref_atom_name_chars",
+            "atom_attention_mask",
+            "token_attention_mask",
+            "residue_index",
+        ):
+            output[name] = features[name]
+        return output
+
+    def fold(
+        self,
+        input,
+        *,
+        num_loops: int = 3,
+        num_sampling_steps: int = 50,
+        num_diffusion_samples: int = 1,
+        seed: int | None = None,
+        noise_scale: float | None = None,
+        step_scale: float | None = None,
+        max_inference_sigma: int | None = None,
+        early_exit: bool = False,
+        complex_id: str = "pred",
+    ):
+        return self.input_builder.fold(
+            self,
+            input,
+            num_loops=num_loops,
+            num_sampling_steps=num_sampling_steps,
+            num_diffusion_samples=num_diffusion_samples,
+            seed=seed,
+            noise_scale=noise_scale,
+            step_scale=step_scale,
+            max_inference_sigma=max_inference_sigma,
+            early_exit=early_exit,
+            complex_id=complex_id,
+        )
+
+    def fold_protein(
+        self,
+        sequence: str,
+        *,
+        chain_id: str = "A",
+        num_loops: int = 3,
+        num_sampling_steps: int = 50,
+        num_diffusion_samples: int = 1,
+        seed: int | None = None,
+        complex_id: str = "pred",
+    ):
+        from .esmfold2_types import ProteinInput, StructurePredictionInput
+
+        input = StructurePredictionInput(
+            sequences=[ProteinInput(id=chain_id, sequence=sequence)]
+        )
+        return self.fold(
+            input,
+            num_loops=num_loops,
+            num_sampling_steps=num_sampling_steps,
+            num_diffusion_samples=num_diffusion_samples,
+            seed=seed,
+            complex_id=complex_id,
+        )
+
+    @staticmethod
+    def result_to_cif(result) -> str:
+        assert not isinstance(result, list), "Pass one MolecularComplexResult at a time."
+        return result.complex.to_mmcif()
+
+    @staticmethod
+    def result_to_pdb(result) -> str:
+        assert not isinstance(result, list), "Pass one MolecularComplexResult at a time."
+        return result.complex.to_protein_complex().to_pdb_string()
+
+    def save_as_cif(self, result, output_path: str | Path) -> None:
+        Path(output_path).write_text(self.result_to_cif(result))
+
+    def save_as_pdb(self, result, output_path: str | Path) -> None:
+        Path(output_path).write_text(self.result_to_pdb(result))
+
+    def infer_protein_as_cif(self, seq: str, **forward_kwargs) -> str:
+        return self.result_to_cif(self.fold_protein(seq, **forward_kwargs))
+
+    def infer_protein_as_pdb(self, seq: str, **forward_kwargs) -> str:
+        return self.result_to_pdb(self.fold_protein(seq, **forward_kwargs))
+
+
+__all__ = [
+    "ConfidenceHead",
+    "MSAEncoder",
+    "MSAEncoderBlock",
+    "ESMFold2ExperimentalModel",
+]

protein_utils.py

@@ -1,488 +1,488 @@
-# coding=utf-8
-# Copyright 2026 Biohub. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Self-contained protein featurization for ESMFold2 inference.
-
-Lets ``ESMFold2ExperimentalModel.infer_protein_as_pdb`` fold a protein sequence
-ESMFold-style without the ``esm`` companion package. The featurization
-mirrors ``ESMFold2InputBuilder.prepare_input`` for the protein-only path —
-``test_prepare_protein_features.py`` enforces tensor-exact parity.
-"""
-
-from __future__ import annotations
-
-import math
-
-import torch
-from torch import Tensor
-
-MOL_TYPE_PROTEIN = 0
-PROTEIN_UNK_RES_TYPE = 22
-MSA_GAP_TOKEN_ID = 1
-
-PROTEIN_RESIDUE_TO_RES_TYPE: dict[str, int] = {
-    "ALA": 2,
-    "ARG": 3,
-    "ASN": 4,
-    "ASP": 5,
-    "CYS": 6,
-    "GLN": 7,
-    "GLU": 8,
-    "GLY": 9,
-    "HIS": 10,
-    "ILE": 11,
-    "LEU": 12,
-    "LYS": 13,
-    "MET": 14,
-    "PHE": 15,
-    "PRO": 16,
-    "SER": 17,
-    "THR": 18,
-    "TRP": 19,
-    "TYR": 20,
-    "VAL": 21,
-}
-
-PROTEIN_1TO3: dict[str, str] = {
-    "A": "ALA",
-    "R": "ARG",
-    "N": "ASN",
-    "D": "ASP",
-    "C": "CYS",
-    "Q": "GLN",
-    "E": "GLU",
-    "G": "GLY",
-    "H": "HIS",
-    "I": "ILE",
-    "L": "LEU",
-    "K": "LYS",
-    "M": "MET",
-    "F": "PHE",
-    "P": "PRO",
-    "S": "SER",
-    "T": "THR",
-    "W": "TRP",
-    "Y": "TYR",
-    "V": "VAL",
-    "X": "UNK",
-}
-
-ESM_PROTEIN_VOCAB: dict[str, int] = {
-    "L": 4,
-    "A": 5,
-    "G": 6,
-    "V": 7,
-    "S": 8,
-    "E": 9,
-    "R": 10,
-    "T": 11,
-    "I": 12,
-    "D": 13,
-    "P": 14,
-    "K": 15,
-    "Q": 16,
-    "N": 17,
-    "F": 18,
-    "Y": 19,
-    "M": 20,
-    "H": 21,
-    "W": 22,
-    "C": 23,
-    "X": 3,
-}
-
-# Heavy atoms per canonical residue, in training-time order.
-PROTEIN_HEAVY_ATOMS: dict[str, list[str]] = {
-    "ALA": ["N", "CA", "C", "O", "CB"],
-    "ARG": ["N", "CA", "C", "O", "CB", "CG", "CD", "NE", "CZ", "NH1", "NH2"],
-    "ASN": ["N", "CA", "C", "O", "CB", "CG", "OD1", "ND2"],
-    "ASP": ["N", "CA", "C", "O", "CB", "CG", "OD1", "OD2"],
-    "CYS": ["N", "CA", "C", "O", "CB", "SG"],
-    "GLN": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "NE2"],
-    "GLU": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "OE2"],
-    "GLY": ["N", "CA", "C", "O"],
-    "HIS": ["N", "CA", "C", "O", "CB", "CG", "ND1", "CD2", "CE1", "NE2"],
-    "ILE": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "CD1"],
-    "LEU": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2"],
-    "LYS": ["N", "CA", "C", "O", "CB", "CG", "CD", "CE", "NZ"],
-    "MET": ["N", "CA", "C", "O", "CB", "CG", "SD", "CE"],
-    "PHE": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ"],
-    "PRO": ["N", "CA", "C", "O", "CB", "CG", "CD"],
-    "SER": ["N", "CA", "C", "O", "CB", "OG"],
-    "THR": ["N", "CA", "C", "O", "CB", "OG1", "CG2"],
-    "TRP": [
-        "N",
-        "CA",
-        "C",
-        "O",
-        "CB",
-        "CG",
-        "CD1",
-        "CD2",
-        "NE1",
-        "CE2",
-        "CE3",
-        "CZ2",
-        "CZ3",
-        "CH2",
-    ],
-    "TYR": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "OH"],
-    "VAL": ["N", "CA", "C", "O", "CB", "CG1", "CG2"],
-    "UNK": ["N", "CA", "C", "O"],
-}
-
-PROTEIN_REF_POS: dict[str, dict[str, tuple[float, float, float]]] = {
-    "ALA": {
-        "N": (-0.01003183238208294, -1.2073018550872803, -1.0555061101913452),
-        "CA": (-0.04190138354897499, 0.17447763681411743, -0.5729365348815918),
-        "C": (1.2127548456192017, 0.4737588167190552, 0.19521640241146088),
-        "O": (1.9390329122543335, 1.4484562873840332, -0.13759790360927582),
-        "CB": (-1.276943325996399, 0.4288230538368225, 0.29937705397605896),
-    },
-    "ARG": {
-        "N": (-2.0170421600341797, 0.6717798113822937, -1.1794233322143555),
-        "CA": (-2.0503084659576416, -0.5735036730766296, -0.4097220301628113),
-        "C": (-3.469440460205078, -1.0612813234329224, -0.2755832374095917),
-        "O": (-3.8218462467193604, -2.1369943618774414, -0.8294969797134399),
-        "CB": (-1.4193516969680786, -0.3735991418361664, 0.9852858781814575),
-        "CG": (0.11878877878189087, -0.3112654983997345, 0.963895857334137),
-        "CD": (0.6643245816230774, 1.0068185329437256, 0.3963329493999481),
-        "NE": (2.1090238094329834, 1.0977025032043457, 0.6120952367782593),
-        "CZ": (3.098905324935913, 0.3215920031070709, -0.09047172218561172),
-        "NH1": (4.461230278015137, 0.3844667971134186, 0.34141138195991516),
-        "NH2": (2.7856509685516357, -0.4166366159915924, -1.1148239374160767),
-    },
-    "ASN": {
-        "N": (-0.7595629096031189, 0.7503494620323181, 1.1369825601577759),
-        "CA": (-0.76087886095047, 0.23876343667507172, -0.23573364317417145),
-        "C": (-1.9211044311523438, -0.6982439160346985, -0.42196929454803467),
-        "O": (-2.677666187286377, -0.5753439664840698, -1.4223182201385498),
-        "CB": (0.5504899024963379, -0.5078350305557251, -0.5390339493751526),
-        "CG": (1.7250099182128906, 0.4264017939567566, -0.5778228640556335),
-        "OD1": (1.9470350742340088, 1.1086392402648926, -1.613560438156128),
-        "ND2": (2.57365345954895, 0.5730618834495544, 0.5608599781990051),
-    },
-    "ASP": {
-        "N": (-1.8452696800231934, -1.2169504165649414, 0.19437327980995178),
-        "CA": (-0.6379959583282471, -0.41974392533302307, 0.41681644320487976),
-        "C": (-0.9431572556495667, 1.0356197357177734, 0.18555717170238495),
-        "O": (-1.5183608531951904, 1.4045922756195068, -0.8739855885505676),
-        "CB": (0.48594576120376587, -0.8970447778701782, -0.5209363698959351),
-        "CG": (1.780342936515808, -0.19918935000896454, -0.2310730367898941),
-        "OD1": (2.5202910900115967, -0.6044584512710571, 0.7049641013145447),
-        "OD2": (2.1454880237579346, 0.9208861589431763, -0.9712985157966614),
-    },
-    "CYS": {
-        "N": (0.0469963513314724, 1.190075159072876, -1.1607273817062378),
-        "CA": (0.11344368755817413, -0.09400428831577301, -0.45952197909355164),
-        "C": (-1.2652032375335693, -0.6832379698753357, -0.3594406247138977),
-        "O": (-1.4631439447402954, -1.8851220607757568, -0.6826791763305664),
-        "CB": (0.6919880509376526, 0.09034398198127747, 0.952482283115387),
-        "SG": (2.4619927406311035, 0.5235707759857178, 0.9020372629165649),
-    },
-    "GLN": {
-        "N": (-2.370004653930664, -0.9637529850006104, -0.7942749261856079),
-        "CA": (-1.370002269744873, -0.6000258922576904, 0.2103111445903778),
-        "C": (-1.7545503377914429, 0.7091967463493347, 0.8433493971824646),
-        "O": (-1.8520662784576416, 0.7999289631843567, 2.0964975357055664),
-        "CB": (0.02040259726345539, -0.5004461407661438, -0.44764479994773865),
-        "CG": (1.1377512216567993, -0.28680720925331116, 0.582992434501648),
-        "CD": (2.4745187759399414, -0.24800164997577667, -0.09364881366491318),
-        "OE1": (3.1685523986816406, -1.2966246604919434, -0.1717153936624527),
-        "NE2": (2.947425603866577, 0.9601329565048218, -0.6888364553451538),
-    },
-    "GLU": {
-        "N": (-1.5850872993469238, -1.337684154510498, 0.9490851163864136),
-        "CA": (-1.0560977458953857, 0.027459044009447098, 1.0306966304779053),
-        "C": (-1.7741456031799316, 0.9664392471313477, 0.09259600937366486),
-        "O": (-1.9012441635131836, 2.181349992752075, 0.402479350566864),
-        "CB": (0.4706551432609558, 0.048803869634866714, 0.8114414811134338),
-        "CG": (0.9133604764938354, -0.4219329059123993, -0.5830985307693481),
-        "CD": (2.398822069168091, -0.3097084164619446, -0.7210537791252136),
-        "OE1": (3.1389315128326416, -1.274524450302124, -0.39029765129089355),
-        "OE2": (2.9647817611694336, 0.8781346082687378, -1.1732689142227173),
-    },
-    "GLY": {
-        "N": (-1.3942985534667969, -0.39875128865242004, -0.3370324671268463),
-        "CA": (-0.39974430203437805, 0.5488945245742798, 0.15242962539196014),
-        "C": (0.9440054893493652, -0.10314033925533295, 0.19859643280506134),
-        "O": (1.3352899551391602, -0.669218122959137, 1.2541258335113525),
-    },
-    "HIS": {
-        "N": (-1.4532867670059204, -1.0689626932144165, 0.881072461605072),
-        "CA": (-1.3396095037460327, 0.24797579646110535, 0.24960045516490936),
-        "C": (-2.675257921218872, 0.6571555733680725, -0.30441102385520935),
-        "O": (-3.1311378479003906, 1.8079776763916016, -0.06785715371370316),
-        "CB": (-0.3041955828666687, 0.21721023321151733, -0.8885309100151062),
-        "CG": (1.0887513160705566, 0.028941065073013306, -0.36419469118118286),
-        "ND1": (1.840459942817688, 1.0411773920059204, 0.29804590344429016),
-        "CD2": (1.780855417251587, -1.1011489629745483, -0.3814258575439453),
-        "CE1": (2.9566943645477295, 0.4924798905849457, 0.6477115750312805),
-        "NE2": (3.0280203819274902, -0.8751969337463379, 0.26084381341934204),
-    },
-    "ILE": {
-        "N": (-0.7167549729347229, -1.5426139831542969, -0.9983330368995667),
-        "CA": (-1.0636085271835327, -0.35169270634651184, -0.21393552422523499),
-        "C": (-1.3896740674972534, 0.8142145276069641, -1.1164065599441528),
-        "O": (-1.2377792596817017, 0.7302915453910828, -2.3656840324401855),
-        "CB": (0.061667006462812424, 0.01599610224366188, 0.8057394623756409),
-        "CG1": (1.502519965171814, -0.08899776637554169, 0.24154816567897797),
-        "CG2": (-0.053174979984760284, -0.8521055579185486, 2.0702083110809326),
-        "CD1": (1.7929610013961792, 0.899773120880127, -0.8863027691841125),
-    },
-    "LEU": {
-        "N": (1.9657520055770874, -1.9763224124908447, -0.18391533195972443),
-        "CA": (1.3077669143676758, -0.6677430868148804, -0.19492436945438385),
-        "C": (1.9905058145523071, 0.24182087182998657, 0.7879968285560608),
-        "O": (2.06896710395813, -0.07880014181137085, 2.0048046112060547),
-        "CB": (-0.20306941866874695, -0.8093230128288269, 0.11243502795696259),
-        "CG": (-0.9916267395019531, 0.5234957337379456, 0.06723011285066605),
-        "CD1": (-2.4228057861328125, 0.29949337244033813, 0.573042094707489),
-        "CD2": (-1.0282856225967407, 1.1250264644622803, -1.346014380455017),
-    },
-    "LYS": {
-        "N": (2.4221372604370117, -0.6473312377929688, 0.6370573043823242),
-        "CA": (2.0314927101135254, 0.2786507308483124, -0.4298512041568756),
-        "C": (2.7168593406677246, 1.595757246017456, -0.20924785733222961),
-        "O": (3.397681713104248, 2.116427421569824, -1.1332510709762573),
-        "CB": (0.5018402934074402, 0.4873858690261841, -0.49062973260879517),
-        "CG": (-0.25062066316604614, -0.7894009947776794, -0.9055535793304443),
-        "CD": (-1.769762635231018, -0.5552700161933899, -1.040329933166504),
-        "CE": (-2.576533555984497, -1.0221366882324219, 0.18493641912937164),
-        "NZ": (-2.269151210784912, -0.24293844401836395, 1.3849012851715088),
-    },
-    "MET": {
-        "N": (1.8903918266296387, -1.5252995491027832, -0.42638593912124634),
-        "CA": (1.2630571126937866, -0.24417810142040253, -0.7626462578773499),
-        "C": (2.30391001701355, 0.8367712497711182, -0.7254616618156433),
-        "O": (2.465414524078369, 1.5928632020950317, -1.7207728624343872),
-        "CB": (0.10567972809076309, 0.10861825942993164, 0.19741646945476532),
-        "CG": (-1.0658042430877686, -0.8736631274223328, 0.08811883628368378),
-        "SD": (-2.4557132720947266, -0.3332225978374481, 1.1461700201034546),
-        "CE": (-3.265165090560913, 0.7033554911613464, -0.11588376015424728),
-    },
-    "PHE": {
-        "N": (-2.8484435081481934, -1.525790810585022, 0.01789816841483116),
-        "CA": (-1.591969609260559, -0.8545162677764893, 0.35214468836784363),
-        "C": (-1.8900631666183472, 0.45833414793014526, 1.0232222080230713),
-        "O": (-1.3424992561340332, 0.74432373046875, 2.121629476547241),
-        "CB": (-0.760358452796936, -0.6342853307723999, -0.9257160425186157),
-        "CG": (0.604112982749939, -0.07200468331575394, -0.6148118376731873),
-        "CD1": (0.8468314409255981, 1.2480632066726685, -0.7146694660186768),
-        "CD2": (1.6827683448791504, -0.9758077263832092, -0.1423054188489914),
-        "CE1": (2.1801748275756836, 1.7875733375549316, -0.3744623064994812),
-        "CE2": (2.888307809829712, -0.48277512192726135, 0.16804970800876617),
-        "CZ": (3.149812936782837, 0.9656873941421509, 0.04440271109342575),
-    },
-    "PRO": {
-        "N": (-0.836250364780426, -0.9899801015853882, 0.5561304688453674),
-        "CA": (0.32722190022468567, -0.6164458394050598, -0.25072571635246277),
-        "C": (1.6121541261672974, -1.1711241006851196, 0.31082412600517273),
-        "O": (1.6127740144729614, -2.2771971225738525, 0.9156193733215332),
-        "CB": (0.3248198926448822, 0.9028244018554688, -0.33368146419525146),
-        "CG": (-1.1425083875656128, 1.2730128765106201, -0.2590600252151489),
-        "CD": (-1.8495968580245972, 0.026575811207294464, 0.2681289613246918),
-    },
-    "SER": {
-        "N": (0.674650251865387, 1.5018702745437622, -0.5367295145988464),
-        "CA": (0.00013792862591799349, 0.4966467022895813, 0.28510504961013794),
-        "C": (0.9941009879112244, -0.5374617576599121, 0.73505038022995),
-        "O": (1.0545241832733154, -0.8683545589447021, 1.9495396614074707),
-        "CB": (-1.1279288530349731, -0.1659376323223114, -0.5160963535308838),
-        "OG": (-1.8135979175567627, -1.085249662399292, 0.28947514295578003),
-    },
-    "THR": {
-        "N": (-1.325830340385437, -1.3728225231170654, 0.6882233023643494),
-        "CA": (-0.5433306097984314, -0.16364754736423492, 0.41697052121162415),
-        "C": (-1.294381856918335, 0.7077372074127197, -0.5549946427345276),
-        "O": (-1.6939635276794434, 0.23654410243034363, -1.6540418863296509),
-        "CB": (0.853203296661377, -0.5363803505897522, -0.14109353721141815),
-        "OG1": (1.5220820903778076, -1.379003643989563, 0.7635167837142944),
-        "CG2": (1.7225933074951172, 0.7054727077484131, -0.3651331067085266),
-    },
-    "TRP": {
-        "N": (3.686030864715576, 0.7599999904632568, 0.496155709028244),
-        "CA": (2.384092092514038, 0.09079249948263168, 0.5325262546539307),
-        "C": (2.1113572120666504, -0.6121063232421875, -0.7733646035194397),
-        "O": (1.796526312828064, -1.8323148488998413, -0.7775964140892029),
-        "CB": (1.281521201133728, 1.1139036417007446, 0.8559791445732117),
-        "CG": (-0.04292375594377518, 0.44645074009895325, 1.0942792892456055),
-        "CD1": (-0.42329534888267517, -0.15470874309539795, 2.2227554321289062),
-        "CD2": (-1.1023900508880615, 0.2158389836549759, 0.11529432237148285),
-        "NE1": (-1.7030320167541504, -0.7665823101997375, 2.0595016479492188),
-        "CE2": (-2.045644998550415, -0.4881173074245453, 0.710669219493866),
-        "CE3": (-1.2173502445220947, 0.6102271676063538, -1.300106406211853),
-        "CZ2": (-3.256009340286255, -0.9164394736289978, -0.00984987337142229),
-        "CZ3": (-2.315925121307373, 0.2306906282901764, -1.9776310920715332),
-        "CH2": (-3.3817875385284424, -0.5677337646484375, -1.3032053709030151),
-    },
-    "TYR": {
-        "N": (-1.7900604009628296, -0.8409399390220642, 1.3180142641067505),
-        "CA": (-1.913882851600647, 0.23552845418453217, 0.330669641494751),
-        "C": (-3.347280740737915, 0.3588399887084961, -0.09830684959888458),
-        "O": (-3.967811346054077, -0.6449354290962219, -0.5423302054405212),
-        "CB": (-1.0093992948532104, 0.0004731413209810853, -0.8981552124023438),
-        "CG": (0.4520410895347595, 0.021162061020731926, -0.5305932760238647),
-        "CD1": (1.0992432832717896, 1.1877919435501099, -0.3579142987728119),
-        "CD2": (1.1803174018859863, -1.253401279449463, -0.31122180819511414),
-        "CE1": (2.5253450870513916, 1.1990256309509277, 0.029804613441228867),
-        "CE2": (2.471151113510132, -1.240687608718872, 0.043534230440855026),
-        "CZ": (3.180687665939331, 0.04672492295503616, 0.2214856892824173),
-        "OH": (4.523719787597656, 0.0671030730009079, 0.5877485871315002),
-    },
-    "VAL": {
-        "N": (0.5987519025802612, -1.569443702697754, -0.7379124760627747),
-        "CA": (0.6014357209205627, -0.10503966361284256, -0.6336286664009094),
-        "C": (1.8391697406768799, 0.4067850410938263, 0.06351757049560547),
-        "O": (2.3952062129974365, -0.2666190266609192, 0.9731166958808899),
-        "CB": (-0.694736897945404, 0.4259096384048462, 0.03581475466489792),
-        "CG1": (-1.9276031255722046, 0.09515828639268875, -0.8172357082366943),
-        "CG2": (-0.8938426971435547, -0.08640842139720917, 1.472349762916565),
-    },
-    "UNK": {
-        "N": (0.0, 0.0, 0.0),
-        "CA": (0.0, 0.0, 0.0),
-        "C": (0.0, 0.0, 0.0),
-        "O": (0.0, 0.0, 0.0),
-    },
-}
-
-# Protonated nitrogens at physiological pH (matches CHARGED_ATOMS in the
-# opensource constants for the protein subset).
-PROTEIN_CHARGED_ATOMS: dict[tuple[str, str], int] = {
-    ("LYS", "NZ"): 1,
-    ("ARG", "NH2"): 1,
-    ("HIS", "ND1"): 1,
-}
-
-# Only the elements that appear in canonical protein heavy atoms.
-_PROTEIN_ELEMENT_TO_ATOMIC_NUM: dict[str, int] = {"C": 6, "N": 7, "O": 8, "S": 16}
-
-
-def _encode_atom_name(name: str) -> list[int]:
-    padded = name.ljust(4)[:4]
-    return [ord(c) - 32 if c != " " else 0 for c in padded]
-
-
-def prepare_protein_features(sequence: str) -> dict[str, Tensor]:
-    """Featurize a single protein sequence for ESMFold2ExperimentalModel.forward.
-
-    Returns the same keys with the same dtypes/shapes as
-    ``ESMFold2InputBuilder.prepare_input(StructurePredictionInput(...))``
-    restricted to a single-chain protein with no MSA, modifications,
-    distogram conditioning, or covalent bonds. All tensors have a
-    leading batch dim of 1; the caller is responsible for moving them
-    to the model device.
-    """
-    if not sequence:
-        raise ValueError("sequence must be non-empty")
-
-    res_3letter = [PROTEIN_1TO3.get(c, "UNK") for c in sequence]
-    L = len(sequence)
-
-    token_atom_starts: list[int] = []
-    atom_records: list[tuple[int, str, str, int, tuple[float, float, float]]] = []
-    res_type_vals: list[int] = []
-    input_id_vals: list[int] = []
-    distogram_rep_atom_idx: list[int] = []
-
-    atom_cursor = 0
-    for t_idx, (letter, res_3) in enumerate(zip(sequence, res_3letter)):
-        atom_names = PROTEIN_HEAVY_ATOMS[res_3]
-        res_type = PROTEIN_RESIDUE_TO_RES_TYPE.get(res_3, PROTEIN_UNK_RES_TYPE)
-        input_id = ESM_PROTEIN_VOCAB.get(letter, ESM_PROTEIN_VOCAB["X"])
-
-        token_atom_starts.append(atom_cursor)
-        for name in atom_names:
-            charge = PROTEIN_CHARGED_ATOMS.get((res_3, name), 0)
-            element = name[0]  # protein heavy atoms are all single-letter C/N/O/S
-            ref_pos = PROTEIN_REF_POS[res_3][name]
-            atom_records.append((t_idx, name, element, charge, ref_pos))
-            atom_cursor += 1
-
-        rep_name = "CB" if "CB" in atom_names else "CA"
-        distogram_rep_atom_idx.append(
-            token_atom_starts[t_idx] + atom_names.index(rep_name)
-        )
-
-        res_type_vals.append(res_type)
-        input_id_vals.append(input_id)
-
-    n_real_atoms = len(atom_records)
-    n_atoms = math.ceil(n_real_atoms / 32) * 32 if n_real_atoms > 0 else 32
-
-    ref_pos = torch.zeros(n_atoms, 3, dtype=torch.float32)
-    ref_element = torch.zeros(n_atoms, dtype=torch.int64)
-    ref_charge = torch.zeros(n_atoms, dtype=torch.int8)
-    ref_atom_name_chars = torch.zeros(n_atoms, 4, dtype=torch.int64)
-    ref_space_uid = torch.zeros(n_atoms, dtype=torch.int64)
-    atom_attention_mask = torch.zeros(n_atoms, dtype=torch.bool)
-    atom_to_token = torch.zeros(n_atoms, dtype=torch.int64)
-
-    for i, (t_idx, name, element, charge, pos) in enumerate(atom_records):
-        ref_pos[i] = torch.tensor(pos, dtype=torch.float32)
-        ref_element[i] = _PROTEIN_ELEMENT_TO_ATOMIC_NUM[element]
-        ref_charge[i] = charge
-        ref_atom_name_chars[i] = torch.tensor(
-            _encode_atom_name(name), dtype=torch.int64
-        )
-        ref_space_uid[i] = t_idx
-        atom_attention_mask[i] = True
-        atom_to_token[i] = t_idx
-
-    token_index = torch.arange(L, dtype=torch.int64)
-    residue_index = torch.arange(L, dtype=torch.int64)
-    asym_id = torch.zeros(L, dtype=torch.int64)
-    sym_id = torch.zeros(L, dtype=torch.int64)
-    entity_id = torch.ones(L, dtype=torch.int64)
-    mol_type = torch.full((L,), MOL_TYPE_PROTEIN, dtype=torch.int64)
-    res_type = torch.tensor(res_type_vals, dtype=torch.int64)
-    input_ids = torch.tensor(input_id_vals, dtype=torch.int64)
-    token_bonds = torch.zeros(L, L, 1, dtype=torch.float32)
-    token_attention_mask = torch.ones(L, dtype=torch.bool)
-    distogram_atom_idx = torch.tensor(distogram_rep_atom_idx, dtype=torch.int64)
-
-    # Single-sequence MSA: depth 1, row 0 is the sequence itself.
-    msa = res_type.unsqueeze(0)
-    msa_attention_mask = torch.ones(1, L, dtype=torch.bool)
-    has_deletion = torch.zeros(1, L, dtype=torch.bool)
-    deletion_value = torch.zeros(1, L, dtype=torch.float32)
-    deletion_mean = torch.zeros(L, dtype=torch.float32)
-
-    features = {
-        "token_index": token_index,
-        "residue_index": residue_index,
-        "asym_id": asym_id,
-        "sym_id": sym_id,
-        "entity_id": entity_id,
-        "mol_type": mol_type,
-        "res_type": res_type,
-        "input_ids": input_ids,
-        "token_bonds": token_bonds,
-        "token_attention_mask": token_attention_mask,
-        "ref_pos": ref_pos,
-        "ref_element": ref_element,
-        "ref_charge": ref_charge,
-        "ref_atom_name_chars": ref_atom_name_chars,
-        "ref_space_uid": ref_space_uid,
-        "atom_attention_mask": atom_attention_mask,
-        "atom_to_token": atom_to_token,
-        "distogram_atom_idx": distogram_atom_idx,
-        "msa": msa,
-        "msa_attention_mask": msa_attention_mask,
-        "has_deletion": has_deletion,
-        "deletion_value": deletion_value,
-        "deletion_mean": deletion_mean,
-    }
-    return {k: v.unsqueeze(0) for k, v in features.items()}
+# coding=utf-8
+# Copyright 2026 Biohub. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Self-contained protein featurization for ESMFold2 inference.
+
+Lets ``ESMFold2ExperimentalModel.infer_protein_as_pdb`` fold a protein sequence
+ESMFold-style without the ``esm`` companion package. The featurization
+mirrors ``ESMFold2InputBuilder.prepare_input`` for the protein-only path —
+``test_prepare_protein_features.py`` enforces tensor-exact parity.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+from torch import Tensor
+
+MOL_TYPE_PROTEIN = 0
+PROTEIN_UNK_RES_TYPE = 22
+MSA_GAP_TOKEN_ID = 1
+
+PROTEIN_RESIDUE_TO_RES_TYPE: dict[str, int] = {
+    "ALA": 2,
+    "ARG": 3,
+    "ASN": 4,
+    "ASP": 5,
+    "CYS": 6,
+    "GLN": 7,
+    "GLU": 8,
+    "GLY": 9,
+    "HIS": 10,
+    "ILE": 11,
+    "LEU": 12,
+    "LYS": 13,
+    "MET": 14,
+    "PHE": 15,
+    "PRO": 16,
+    "SER": 17,
+    "THR": 18,
+    "TRP": 19,
+    "TYR": 20,
+    "VAL": 21,
+}
+
+PROTEIN_1TO3: dict[str, str] = {
+    "A": "ALA",
+    "R": "ARG",
+    "N": "ASN",
+    "D": "ASP",
+    "C": "CYS",
+    "Q": "GLN",
+    "E": "GLU",
+    "G": "GLY",
+    "H": "HIS",
+    "I": "ILE",
+    "L": "LEU",
+    "K": "LYS",
+    "M": "MET",
+    "F": "PHE",
+    "P": "PRO",
+    "S": "SER",
+    "T": "THR",
+    "W": "TRP",
+    "Y": "TYR",
+    "V": "VAL",
+    "X": "UNK",
+}
+
+ESM_PROTEIN_VOCAB: dict[str, int] = {
+    "L": 4,
+    "A": 5,
+    "G": 6,
+    "V": 7,
+    "S": 8,
+    "E": 9,
+    "R": 10,
+    "T": 11,
+    "I": 12,
+    "D": 13,
+    "P": 14,
+    "K": 15,
+    "Q": 16,
+    "N": 17,
+    "F": 18,
+    "Y": 19,
+    "M": 20,
+    "H": 21,
+    "W": 22,
+    "C": 23,
+    "X": 3,
+}
+
+# Heavy atoms per canonical residue, in training-time order.
+PROTEIN_HEAVY_ATOMS: dict[str, list[str]] = {
+    "ALA": ["N", "CA", "C", "O", "CB"],
+    "ARG": ["N", "CA", "C", "O", "CB", "CG", "CD", "NE", "CZ", "NH1", "NH2"],
+    "ASN": ["N", "CA", "C", "O", "CB", "CG", "OD1", "ND2"],
+    "ASP": ["N", "CA", "C", "O", "CB", "CG", "OD1", "OD2"],
+    "CYS": ["N", "CA", "C", "O", "CB", "SG"],
+    "GLN": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "NE2"],
+    "GLU": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "OE2"],
+    "GLY": ["N", "CA", "C", "O"],
+    "HIS": ["N", "CA", "C", "O", "CB", "CG", "ND1", "CD2", "CE1", "NE2"],
+    "ILE": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "CD1"],
+    "LEU": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2"],
+    "LYS": ["N", "CA", "C", "O", "CB", "CG", "CD", "CE", "NZ"],
+    "MET": ["N", "CA", "C", "O", "CB", "CG", "SD", "CE"],
+    "PHE": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ"],
+    "PRO": ["N", "CA", "C", "O", "CB", "CG", "CD"],
+    "SER": ["N", "CA", "C", "O", "CB", "OG"],
+    "THR": ["N", "CA", "C", "O", "CB", "OG1", "CG2"],
+    "TRP": [
+        "N",
+        "CA",
+        "C",
+        "O",
+        "CB",
+        "CG",
+        "CD1",
+        "CD2",
+        "NE1",
+        "CE2",
+        "CE3",
+        "CZ2",
+        "CZ3",
+        "CH2",
+    ],
+    "TYR": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "OH"],
+    "VAL": ["N", "CA", "C", "O", "CB", "CG1", "CG2"],
+    "UNK": ["N", "CA", "C", "O"],
+}
+
+PROTEIN_REF_POS: dict[str, dict[str, tuple[float, float, float]]] = {
+    "ALA": {
+        "N": (-0.01003183238208294, -1.2073018550872803, -1.0555061101913452),
+        "CA": (-0.04190138354897499, 0.17447763681411743, -0.5729365348815918),
+        "C": (1.2127548456192017, 0.4737588167190552, 0.19521640241146088),
+        "O": (1.9390329122543335, 1.4484562873840332, -0.13759790360927582),
+        "CB": (-1.276943325996399, 0.4288230538368225, 0.29937705397605896),
+    },
+    "ARG": {
+        "N": (-2.0170421600341797, 0.6717798113822937, -1.1794233322143555),
+        "CA": (-2.0503084659576416, -0.5735036730766296, -0.4097220301628113),
+        "C": (-3.469440460205078, -1.0612813234329224, -0.2755832374095917),
+        "O": (-3.8218462467193604, -2.1369943618774414, -0.8294969797134399),
+        "CB": (-1.4193516969680786, -0.3735991418361664, 0.9852858781814575),
+        "CG": (0.11878877878189087, -0.3112654983997345, 0.963895857334137),
+        "CD": (0.6643245816230774, 1.0068185329437256, 0.3963329493999481),
+        "NE": (2.1090238094329834, 1.0977025032043457, 0.6120952367782593),
+        "CZ": (3.098905324935913, 0.3215920031070709, -0.09047172218561172),
+        "NH1": (4.461230278015137, 0.3844667971134186, 0.34141138195991516),
+        "NH2": (2.7856509685516357, -0.4166366159915924, -1.1148239374160767),
+    },
+    "ASN": {
+        "N": (-0.7595629096031189, 0.7503494620323181, 1.1369825601577759),
+        "CA": (-0.76087886095047, 0.23876343667507172, -0.23573364317417145),
+        "C": (-1.9211044311523438, -0.6982439160346985, -0.42196929454803467),
+        "O": (-2.677666187286377, -0.5753439664840698, -1.4223182201385498),
+        "CB": (0.5504899024963379, -0.5078350305557251, -0.5390339493751526),
+        "CG": (1.7250099182128906, 0.4264017939567566, -0.5778228640556335),
+        "OD1": (1.9470350742340088, 1.1086392402648926, -1.613560438156128),
+        "ND2": (2.57365345954895, 0.5730618834495544, 0.5608599781990051),
+    },
+    "ASP": {
+        "N": (-1.8452696800231934, -1.2169504165649414, 0.19437327980995178),
+        "CA": (-0.6379959583282471, -0.41974392533302307, 0.41681644320487976),
+        "C": (-0.9431572556495667, 1.0356197357177734, 0.18555717170238495),
+        "O": (-1.5183608531951904, 1.4045922756195068, -0.8739855885505676),
+        "CB": (0.48594576120376587, -0.8970447778701782, -0.5209363698959351),
+        "CG": (1.780342936515808, -0.19918935000896454, -0.2310730367898941),
+        "OD1": (2.5202910900115967, -0.6044584512710571, 0.7049641013145447),
+        "OD2": (2.1454880237579346, 0.9208861589431763, -0.9712985157966614),
+    },
+    "CYS": {
+        "N": (0.0469963513314724, 1.190075159072876, -1.1607273817062378),
+        "CA": (0.11344368755817413, -0.09400428831577301, -0.45952197909355164),
+        "C": (-1.2652032375335693, -0.6832379698753357, -0.3594406247138977),
+        "O": (-1.4631439447402954, -1.8851220607757568, -0.6826791763305664),
+        "CB": (0.6919880509376526, 0.09034398198127747, 0.952482283115387),
+        "SG": (2.4619927406311035, 0.5235707759857178, 0.9020372629165649),
+    },
+    "GLN": {
+        "N": (-2.370004653930664, -0.9637529850006104, -0.7942749261856079),
+        "CA": (-1.370002269744873, -0.6000258922576904, 0.2103111445903778),
+        "C": (-1.7545503377914429, 0.7091967463493347, 0.8433493971824646),
+        "O": (-1.8520662784576416, 0.7999289631843567, 2.0964975357055664),
+        "CB": (0.02040259726345539, -0.5004461407661438, -0.44764479994773865),
+        "CG": (1.1377512216567993, -0.28680720925331116, 0.582992434501648),
+        "CD": (2.4745187759399414, -0.24800164997577667, -0.09364881366491318),
+        "OE1": (3.1685523986816406, -1.2966246604919434, -0.1717153936624527),
+        "NE2": (2.947425603866577, 0.9601329565048218, -0.6888364553451538),
+    },
+    "GLU": {
+        "N": (-1.5850872993469238, -1.337684154510498, 0.9490851163864136),
+        "CA": (-1.0560977458953857, 0.027459044009447098, 1.0306966304779053),
+        "C": (-1.7741456031799316, 0.9664392471313477, 0.09259600937366486),
+        "O": (-1.9012441635131836, 2.181349992752075, 0.402479350566864),
+        "CB": (0.4706551432609558, 0.048803869634866714, 0.8114414811134338),
+        "CG": (0.9133604764938354, -0.4219329059123993, -0.5830985307693481),
+        "CD": (2.398822069168091, -0.3097084164619446, -0.7210537791252136),
+        "OE1": (3.1389315128326416, -1.274524450302124, -0.39029765129089355),
+        "OE2": (2.9647817611694336, 0.8781346082687378, -1.1732689142227173),
+    },
+    "GLY": {
+        "N": (-1.3942985534667969, -0.39875128865242004, -0.3370324671268463),
+        "CA": (-0.39974430203437805, 0.5488945245742798, 0.15242962539196014),
+        "C": (0.9440054893493652, -0.10314033925533295, 0.19859643280506134),
+        "O": (1.3352899551391602, -0.669218122959137, 1.2541258335113525),
+    },
+    "HIS": {
+        "N": (-1.4532867670059204, -1.0689626932144165, 0.881072461605072),
+        "CA": (-1.3396095037460327, 0.24797579646110535, 0.24960045516490936),
+        "C": (-2.675257921218872, 0.6571555733680725, -0.30441102385520935),
+        "O": (-3.1311378479003906, 1.8079776763916016, -0.06785715371370316),
+        "CB": (-0.3041955828666687, 0.21721023321151733, -0.8885309100151062),
+        "CG": (1.0887513160705566, 0.028941065073013306, -0.36419469118118286),
+        "ND1": (1.840459942817688, 1.0411773920059204, 0.29804590344429016),
+        "CD2": (1.780855417251587, -1.1011489629745483, -0.3814258575439453),
+        "CE1": (2.9566943645477295, 0.4924798905849457, 0.6477115750312805),
+        "NE2": (3.0280203819274902, -0.8751969337463379, 0.26084381341934204),
+    },
+    "ILE": {
+        "N": (-0.7167549729347229, -1.5426139831542969, -0.9983330368995667),
+        "CA": (-1.0636085271835327, -0.35169270634651184, -0.21393552422523499),
+        "C": (-1.3896740674972534, 0.8142145276069641, -1.1164065599441528),
+        "O": (-1.2377792596817017, 0.7302915453910828, -2.3656840324401855),
+        "CB": (0.061667006462812424, 0.01599610224366188, 0.8057394623756409),
+        "CG1": (1.502519965171814, -0.08899776637554169, 0.24154816567897797),
+        "CG2": (-0.053174979984760284, -0.8521055579185486, 2.0702083110809326),
+        "CD1": (1.7929610013961792, 0.899773120880127, -0.8863027691841125),
+    },
+    "LEU": {
+        "N": (1.9657520055770874, -1.9763224124908447, -0.18391533195972443),
+        "CA": (1.3077669143676758, -0.6677430868148804, -0.19492436945438385),
+        "C": (1.9905058145523071, 0.24182087182998657, 0.7879968285560608),
+        "O": (2.06896710395813, -0.07880014181137085, 2.0048046112060547),
+        "CB": (-0.20306941866874695, -0.8093230128288269, 0.11243502795696259),
+        "CG": (-0.9916267395019531, 0.5234957337379456, 0.06723011285066605),
+        "CD1": (-2.4228057861328125, 0.29949337244033813, 0.573042094707489),
+        "CD2": (-1.0282856225967407, 1.1250264644622803, -1.346014380455017),
+    },
+    "LYS": {
+        "N": (2.4221372604370117, -0.6473312377929688, 0.6370573043823242),
+        "CA": (2.0314927101135254, 0.2786507308483124, -0.4298512041568756),
+        "C": (2.7168593406677246, 1.595757246017456, -0.20924785733222961),
+        "O": (3.397681713104248, 2.116427421569824, -1.1332510709762573),
+        "CB": (0.5018402934074402, 0.4873858690261841, -0.49062973260879517),
+        "CG": (-0.25062066316604614, -0.7894009947776794, -0.9055535793304443),
+        "CD": (-1.769762635231018, -0.5552700161933899, -1.040329933166504),
+        "CE": (-2.576533555984497, -1.0221366882324219, 0.18493641912937164),
+        "NZ": (-2.269151210784912, -0.24293844401836395, 1.3849012851715088),
+    },
+    "MET": {
+        "N": (1.8903918266296387, -1.5252995491027832, -0.42638593912124634),
+        "CA": (1.2630571126937866, -0.24417810142040253, -0.7626462578773499),
+        "C": (2.30391001701355, 0.8367712497711182, -0.7254616618156433),
+        "O": (2.465414524078369, 1.5928632020950317, -1.7207728624343872),
+        "CB": (0.10567972809076309, 0.10861825942993164, 0.19741646945476532),
+        "CG": (-1.0658042430877686, -0.8736631274223328, 0.08811883628368378),
+        "SD": (-2.4557132720947266, -0.3332225978374481, 1.1461700201034546),
+        "CE": (-3.265165090560913, 0.7033554911613464, -0.11588376015424728),
+    },
+    "PHE": {
+        "N": (-2.8484435081481934, -1.525790810585022, 0.01789816841483116),
+        "CA": (-1.591969609260559, -0.8545162677764893, 0.35214468836784363),
+        "C": (-1.8900631666183472, 0.45833414793014526, 1.0232222080230713),
+        "O": (-1.3424992561340332, 0.74432373046875, 2.121629476547241),
+        "CB": (-0.760358452796936, -0.6342853307723999, -0.9257160425186157),
+        "CG": (0.604112982749939, -0.07200468331575394, -0.6148118376731873),
+        "CD1": (0.8468314409255981, 1.2480632066726685, -0.7146694660186768),
+        "CD2": (1.6827683448791504, -0.9758077263832092, -0.1423054188489914),
+        "CE1": (2.1801748275756836, 1.7875733375549316, -0.3744623064994812),
+        "CE2": (2.888307809829712, -0.48277512192726135, 0.16804970800876617),
+        "CZ": (3.149812936782837, 0.9656873941421509, 0.04440271109342575),
+    },
+    "PRO": {
+        "N": (-0.836250364780426, -0.9899801015853882, 0.5561304688453674),
+        "CA": (0.32722190022468567, -0.6164458394050598, -0.25072571635246277),
+        "C": (1.6121541261672974, -1.1711241006851196, 0.31082412600517273),
+        "O": (1.6127740144729614, -2.2771971225738525, 0.9156193733215332),
+        "CB": (0.3248198926448822, 0.9028244018554688, -0.33368146419525146),
+        "CG": (-1.1425083875656128, 1.2730128765106201, -0.2590600252151489),
+        "CD": (-1.8495968580245972, 0.026575811207294464, 0.2681289613246918),
+    },
+    "SER": {
+        "N": (0.674650251865387, 1.5018702745437622, -0.5367295145988464),
+        "CA": (0.00013792862591799349, 0.4966467022895813, 0.28510504961013794),
+        "C": (0.9941009879112244, -0.5374617576599121, 0.73505038022995),
+        "O": (1.0545241832733154, -0.8683545589447021, 1.9495396614074707),
+        "CB": (-1.1279288530349731, -0.1659376323223114, -0.5160963535308838),
+        "OG": (-1.8135979175567627, -1.085249662399292, 0.28947514295578003),
+    },
+    "THR": {
+        "N": (-1.325830340385437, -1.3728225231170654, 0.6882233023643494),
+        "CA": (-0.5433306097984314, -0.16364754736423492, 0.41697052121162415),
+        "C": (-1.294381856918335, 0.7077372074127197, -0.5549946427345276),
+        "O": (-1.6939635276794434, 0.23654410243034363, -1.6540418863296509),
+        "CB": (0.853203296661377, -0.5363803505897522, -0.14109353721141815),
+        "OG1": (1.5220820903778076, -1.379003643989563, 0.7635167837142944),
+        "CG2": (1.7225933074951172, 0.7054727077484131, -0.3651331067085266),
+    },
+    "TRP": {
+        "N": (3.686030864715576, 0.7599999904632568, 0.496155709028244),
+        "CA": (2.384092092514038, 0.09079249948263168, 0.5325262546539307),
+        "C": (2.1113572120666504, -0.6121063232421875, -0.7733646035194397),
+        "O": (1.796526312828064, -1.8323148488998413, -0.7775964140892029),
+        "CB": (1.281521201133728, 1.1139036417007446, 0.8559791445732117),
+        "CG": (-0.04292375594377518, 0.44645074009895325, 1.0942792892456055),
+        "CD1": (-0.42329534888267517, -0.15470874309539795, 2.2227554321289062),
+        "CD2": (-1.1023900508880615, 0.2158389836549759, 0.11529432237148285),
+        "NE1": (-1.7030320167541504, -0.7665823101997375, 2.0595016479492188),
+        "CE2": (-2.045644998550415, -0.4881173074245453, 0.710669219493866),
+        "CE3": (-1.2173502445220947, 0.6102271676063538, -1.300106406211853),
+        "CZ2": (-3.256009340286255, -0.9164394736289978, -0.00984987337142229),
+        "CZ3": (-2.315925121307373, 0.2306906282901764, -1.9776310920715332),
+        "CH2": (-3.3817875385284424, -0.5677337646484375, -1.3032053709030151),
+    },
+    "TYR": {
+        "N": (-1.7900604009628296, -0.8409399390220642, 1.3180142641067505),
+        "CA": (-1.913882851600647, 0.23552845418453217, 0.330669641494751),
+        "C": (-3.347280740737915, 0.3588399887084961, -0.09830684959888458),
+        "O": (-3.967811346054077, -0.6449354290962219, -0.5423302054405212),
+        "CB": (-1.0093992948532104, 0.0004731413209810853, -0.8981552124023438),
+        "CG": (0.4520410895347595, 0.021162061020731926, -0.5305932760238647),
+        "CD1": (1.0992432832717896, 1.1877919435501099, -0.3579142987728119),
+        "CD2": (1.1803174018859863, -1.253401279449463, -0.31122180819511414),
+        "CE1": (2.5253450870513916, 1.1990256309509277, 0.029804613441228867),
+        "CE2": (2.471151113510132, -1.240687608718872, 0.043534230440855026),
+        "CZ": (3.180687665939331, 0.04672492295503616, 0.2214856892824173),
+        "OH": (4.523719787597656, 0.0671030730009079, 0.5877485871315002),
+    },
+    "VAL": {
+        "N": (0.5987519025802612, -1.569443702697754, -0.7379124760627747),
+        "CA": (0.6014357209205627, -0.10503966361284256, -0.6336286664009094),
+        "C": (1.8391697406768799, 0.4067850410938263, 0.06351757049560547),
+        "O": (2.3952062129974365, -0.2666190266609192, 0.9731166958808899),
+        "CB": (-0.694736897945404, 0.4259096384048462, 0.03581475466489792),
+        "CG1": (-1.9276031255722046, 0.09515828639268875, -0.8172357082366943),
+        "CG2": (-0.8938426971435547, -0.08640842139720917, 1.472349762916565),
+    },
+    "UNK": {
+        "N": (0.0, 0.0, 0.0),
+        "CA": (0.0, 0.0, 0.0),
+        "C": (0.0, 0.0, 0.0),
+        "O": (0.0, 0.0, 0.0),
+    },
+}
+
+# Protonated nitrogens at physiological pH (matches CHARGED_ATOMS in the
+# opensource constants for the protein subset).
+PROTEIN_CHARGED_ATOMS: dict[tuple[str, str], int] = {
+    ("LYS", "NZ"): 1,
+    ("ARG", "NH2"): 1,
+    ("HIS", "ND1"): 1,
+}
+
+# Only the elements that appear in canonical protein heavy atoms.
+_PROTEIN_ELEMENT_TO_ATOMIC_NUM: dict[str, int] = {"C": 6, "N": 7, "O": 8, "S": 16}
+
+
+def _encode_atom_name(name: str) -> list[int]:
+    padded = name.ljust(4)[:4]
+    return [ord(c) - 32 if c != " " else 0 for c in padded]
+
+
+def prepare_protein_features(sequence: str) -> dict[str, Tensor]:
+    """Featurize a single protein sequence for ESMFold2ExperimentalModel.forward.
+
+    Returns the same keys with the same dtypes/shapes as
+    ``ESMFold2InputBuilder.prepare_input(StructurePredictionInput(...))``
+    restricted to a single-chain protein with no MSA, modifications,
+    distogram conditioning, or covalent bonds. All tensors have a
+    leading batch dim of 1; the caller is responsible for moving them
+    to the model device.
+    """
+    if not sequence:
+        raise ValueError("sequence must be non-empty")
+
+    res_3letter = [PROTEIN_1TO3.get(c, "UNK") for c in sequence]
+    L = len(sequence)
+
+    token_atom_starts: list[int] = []
+    atom_records: list[tuple[int, str, str, int, tuple[float, float, float]]] = []
+    res_type_vals: list[int] = []
+    input_id_vals: list[int] = []
+    distogram_rep_atom_idx: list[int] = []
+
+    atom_cursor = 0
+    for t_idx, (letter, res_3) in enumerate(zip(sequence, res_3letter)):
+        atom_names = PROTEIN_HEAVY_ATOMS[res_3]
+        res_type = PROTEIN_RESIDUE_TO_RES_TYPE.get(res_3, PROTEIN_UNK_RES_TYPE)
+        input_id = ESM_PROTEIN_VOCAB.get(letter, ESM_PROTEIN_VOCAB["X"])
+
+        token_atom_starts.append(atom_cursor)
+        for name in atom_names:
+            charge = PROTEIN_CHARGED_ATOMS.get((res_3, name), 0)
+            element = name[0]  # protein heavy atoms are all single-letter C/N/O/S
+            ref_pos = PROTEIN_REF_POS[res_3][name]
+            atom_records.append((t_idx, name, element, charge, ref_pos))
+            atom_cursor += 1
+
+        rep_name = "CB" if "CB" in atom_names else "CA"
+        distogram_rep_atom_idx.append(
+            token_atom_starts[t_idx] + atom_names.index(rep_name)
+        )
+
+        res_type_vals.append(res_type)
+        input_id_vals.append(input_id)
+
+    n_real_atoms = len(atom_records)
+    n_atoms = math.ceil(n_real_atoms / 32) * 32 if n_real_atoms > 0 else 32
+
+    ref_pos = torch.zeros(n_atoms, 3, dtype=torch.float32)
+    ref_element = torch.zeros(n_atoms, dtype=torch.int64)
+    ref_charge = torch.zeros(n_atoms, dtype=torch.int8)
+    ref_atom_name_chars = torch.zeros(n_atoms, 4, dtype=torch.int64)
+    ref_space_uid = torch.zeros(n_atoms, dtype=torch.int64)
+    atom_attention_mask = torch.zeros(n_atoms, dtype=torch.bool)
+    atom_to_token = torch.zeros(n_atoms, dtype=torch.int64)
+
+    for i, (t_idx, name, element, charge, pos) in enumerate(atom_records):
+        ref_pos[i] = torch.tensor(pos, dtype=torch.float32)
+        ref_element[i] = _PROTEIN_ELEMENT_TO_ATOMIC_NUM[element]
+        ref_charge[i] = charge
+        ref_atom_name_chars[i] = torch.tensor(
+            _encode_atom_name(name), dtype=torch.int64
+        )
+        ref_space_uid[i] = t_idx
+        atom_attention_mask[i] = True
+        atom_to_token[i] = t_idx
+
+    token_index = torch.arange(L, dtype=torch.int64)
+    residue_index = torch.arange(L, dtype=torch.int64)
+    asym_id = torch.zeros(L, dtype=torch.int64)
+    sym_id = torch.zeros(L, dtype=torch.int64)
+    entity_id = torch.ones(L, dtype=torch.int64)
+    mol_type = torch.full((L,), MOL_TYPE_PROTEIN, dtype=torch.int64)
+    res_type = torch.tensor(res_type_vals, dtype=torch.int64)
+    input_ids = torch.tensor(input_id_vals, dtype=torch.int64)
+    token_bonds = torch.zeros(L, L, 1, dtype=torch.float32)
+    token_attention_mask = torch.ones(L, dtype=torch.bool)
+    distogram_atom_idx = torch.tensor(distogram_rep_atom_idx, dtype=torch.int64)
+
+    # Single-sequence MSA: depth 1, row 0 is the sequence itself.
+    msa = res_type.unsqueeze(0)
+    msa_attention_mask = torch.ones(1, L, dtype=torch.bool)
+    has_deletion = torch.zeros(1, L, dtype=torch.bool)
+    deletion_value = torch.zeros(1, L, dtype=torch.float32)
+    deletion_mean = torch.zeros(L, dtype=torch.float32)
+
+    features = {
+        "token_index": token_index,
+        "residue_index": residue_index,
+        "asym_id": asym_id,
+        "sym_id": sym_id,
+        "entity_id": entity_id,
+        "mol_type": mol_type,
+        "res_type": res_type,
+        "input_ids": input_ids,
+        "token_bonds": token_bonds,
+        "token_attention_mask": token_attention_mask,
+        "ref_pos": ref_pos,
+        "ref_element": ref_element,
+        "ref_charge": ref_charge,
+        "ref_atom_name_chars": ref_atom_name_chars,
+        "ref_space_uid": ref_space_uid,
+        "atom_attention_mask": atom_attention_mask,
+        "atom_to_token": atom_to_token,
+        "distogram_atom_idx": distogram_atom_idx,
+        "msa": msa,
+        "msa_attention_mask": msa_attention_mask,
+        "has_deletion": has_deletion,
+        "deletion_value": deletion_value,
+        "deletion_mean": deletion_mean,
+    }
+    return {k: v.unsqueeze(0) for k, v in features.items()}