Fix ESMFold2 prepared feature forwarding

Add-only FastPLMs files-only publication. Checkpoint weights and complete-artifact attestations are unchanged.

README.md

@@ -152,13 +152,15 @@ complex_result = model.fold(
 print(complex_result.ptm, complex_result.plddt.mean().item())
 ```
 
-The typed interface also supports RNA, modifications, covalent bonds, and
-distogram conditioning. Protein MSA inputs are not supported by this Fast
-checkpoint; every protein chain must use `msa=None`. The public schema recognizes
-`PocketConditioning`, but the pinned official runtime discards it and hard-codes
-a zero pocket feature. FastPLMs therefore rejects non-null pocket conditioning
-instead of silently ignoring it. Prepared `ref_pos` values are component
-reference geometries created during featurization, not target coordinates.
+The typed interface also supports RNA, modifications, and covalent bonds.
+Protein MSA inputs are not supported by this Fast checkpoint; every protein
+chain must use `msa=None`. The public schema recognizes `PocketConditioning` and
+`DistogramConditioning`, but the pinned official forward consumes neither. Its
+feature builder hard-codes a zero pocket feature and constructs distogram tensors
+that the released model ignores. FastPLMs therefore rejects non-null pocket and
+distogram conditioning instead of silently ignoring scientific inputs. Prepared
+`ref_pos` values are component reference geometries created during featurization,
+not target coordinates.
 Predicted coordinates and confidence scores are outputs and do not establish
 biochemical activity.
 

config.json

@@ -29,11 +29,11 @@
   "fastplms_checkpoint_repo_id": "Synthyra/ESMFold2-Fast",
   "fastplms_checkpoint_revision": "407875bfcaa42552bfcb25acd67ee1888b790170",
   "fastplms_model_id": "esmfold2_fast",
-  "fastplms_release_tool_revision": "e6dd397a9ad368c998d714f6bd64d40b533d1ed1",
-  "fastplms_release_tool_sha256": "6d335c05aa49a232086a816deb25d248d1490529e5783acc3355b9bc6f03e0c2",
-  "fastplms_runtime_bundle_sha256": "4646da0fba1e0aa8302ca3e78bbf9f00d1123144940f4f377010af0b21ab86b8",
-  "fastplms_runtime_revision": "e6dd397a9ad368c998d714f6bd64d40b533d1ed1",
-  "fastplms_source_tree_sha256": "332058883bd49a8bfda235a1d983bef21f36de1987e44c1f05ffd4247d0d67f1",
+  "fastplms_release_tool_revision": "b5669feacd99580cab27c9e969092fe9c647b9c6",
+  "fastplms_release_tool_sha256": "19f8b87151866ea1bec0274bb99c98890bfbb70887276e57d3990a96d5c6b553",
+  "fastplms_runtime_bundle_sha256": "6e81dad15aa6df4f53633997daf8747aa49a4afbd030a65937a6f3a7d8159273",
+  "fastplms_runtime_revision": "b5669feacd99580cab27c9e969092fe9c647b9c6",
+  "fastplms_source_tree_sha256": "db6fe0054a134e77d111381b520917ffd0d795e4be0bc2cd0e97d43c59b5f3ee",
   "fastplms_weights_revision": "407875bfcaa42552bfcb25acd67ee1888b790170",
   "folding_trunk": {
     "dropout": 0.25,

fastplms/models/esmfold2/esmfold2_processor.py

@@ -123,6 +123,12 @@ def clean_esmfold2_input(input: StructurePredictionInput) -> StructurePrediction
             "the published ESMFold2 feature pipeline drops it. FastPLMs refuses this "
             "input instead of silently emitting an all-zero pocket feature."
         )
+    if input.distogram_conditioning is not None:
+        raise NotImplementedError(
+            "ESMFold2 distogram conditioning is present in the upstream input schema but "
+            "the published ESMFold2 forward does not consume it. FastPLMs refuses this "
+            "input instead of silently ignoring the supplied distogram."
+        )
 
     cleaned: list[Any] = []
     for item in input.sequences:

fastplms/models/esmfold2/modeling_esmfold2.py

@@ -74,6 +74,7 @@ from .modeling_esmfold2_common import (
     maybe_subsample_msa,
     validate_kernel_backend,
     validate_msa_conditioning_inputs,
+    validate_prepared_auxiliary_inputs,
 )
 
 _ESMC_FP8_LINEAR_SUFFIX = ".attn.out_proj"
@@ -1417,6 +1418,12 @@ class ESMFold2Model(
         output_attentions: bool | None = None,
         output_hidden_states: bool | None = None,
         return_dict: bool | None = None,
+        pocket_feature: Tensor | None = None,
+        gt_coords: Tensor | None = None,
+        is_resolved: Tensor | None = None,
+        frames_idx: Tensor | None = None,
+        disto_cond: Tensor | None = None,
+        disto_cond_mask: Tensor | None = None,
     ) -> ESMFold2Output | tuple[Any, ...]:
         output_hidden_states, return_dict = _resolve_structure_output_controls(
             self.config,
@@ -1432,6 +1439,12 @@ class ESMFold2Model(
             deletion_value=deletion_value,
             deletion_mean=deletion_mean,
         )
+        validate_prepared_auxiliary_inputs(
+            pocket_feature=pocket_feature,
+            disto_cond=disto_cond,
+            disto_cond_mask=disto_cond_mask,
+        )
+        del gt_coords, is_resolved, frames_idx
         tok_mask = token_attention_mask
         atm_mask = atom_attention_mask
         disto_idx = distogram_atom_idx

fastplms/models/esmfold2/modeling_esmfold2_common.py

@@ -57,6 +57,14 @@ MSA_CONDITIONING_INPUT_NAMES = (
     "deletion_value",
     "deletion_mean",
 )
+PREPARED_AUXILIARY_INPUT_NAMES = (
+    "pocket_feature",
+    "gt_coords",
+    "is_resolved",
+    "frames_idx",
+    "disto_cond",
+    "disto_cond_mask",
+)
 
 
 def validate_kernel_backend(backend: str | None) -> None:
@@ -105,6 +113,29 @@ def validate_msa_conditioning_inputs(
         )
 
 
+def validate_prepared_auxiliary_inputs(
+    *,
+    pocket_feature: Tensor | None,
+    disto_cond: Tensor | None,
+    disto_cond_mask: Tensor | None,
+) -> None:
+    """Accept inert upstream features while rejecting unsupported conditioning."""
+
+    if pocket_feature is not None and torch.any(pocket_feature != 0).item():
+        raise NotImplementedError(
+            "The published ESMFold2 forward does not consume pocket conditioning; "
+            "nonzero pocket_feature values are unsupported."
+        )
+    distogram_is_active = (disto_cond is not None and torch.any(disto_cond != 0).item()) or (
+        disto_cond_mask is not None and torch.any(disto_cond_mask).item()
+    )
+    if distogram_is_active:
+        raise NotImplementedError(
+            "The published ESMFold2 forward does not consume distogram conditioning; "
+            "nonzero disto_cond or disto_cond_mask values are unsupported."
+        )
+
+
 def _fused_active(module: nn.Module, tensor: Tensor) -> bool:
     """Return whether an optional fused implementation can handle this call."""
     return (

fastplms/models/esmfold2/modeling_esmfold2_experimental.py

@@ -59,6 +59,7 @@ from .modeling_esmfold2_common import (
     gather_token_to_atom,
     validate_kernel_backend,
     validate_msa_conditioning_inputs,
+    validate_prepared_auxiliary_inputs,
 )
 
 _EPS = 1e-5
@@ -689,6 +690,12 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
         output_attentions: bool | None = None,
         output_hidden_states: bool | None = None,
         return_dict: bool | None = None,
+        pocket_feature: Tensor | None = None,
+        gt_coords: Tensor | None = None,
+        is_resolved: Tensor | None = None,
+        frames_idx: Tensor | None = None,
+        disto_cond: Tensor | None = None,
+        disto_cond_mask: Tensor | None = None,
     ) -> ESMFold2Output | tuple[Any, ...]:
         output_hidden_states, return_dict = _resolve_structure_output_controls(
             self.config,
@@ -704,6 +711,12 @@ class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin,
             deletion_value=deletion_value,
             deletion_mean=deletion_mean,
         )
+        validate_prepared_auxiliary_inputs(
+            pocket_feature=pocket_feature,
+            disto_cond=disto_cond,
+            disto_cond_mask=disto_cond_mask,
+        )
+        del gt_coords, is_resolved, frames_idx
         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

modeling_fastplms.py

@@ -12,7 +12,7 @@ from zipfile import ZIP_DEFLATED, ZipFile
 
 from .fastplms_bundle import RUNTIME_DATA, RUNTIME_HASH
 
-if RUNTIME_HASH != "4646da0fba1e0aa8302ca3e78bbf9f00d1123144940f4f377010af0b21ab86b8":
+if RUNTIME_HASH != "6e81dad15aa6df4f53633997daf8747aa49a4afbd030a65937a6f3a7d8159273":
     raise RuntimeError("FastPLMs runtime identity differs from the bridge.")
 
 _RUNTIME_TEMPORARIES = []

runtime-attestation.json

@@ -6,9 +6,9 @@
     "LICENSES/biohub-transformers/LICENSE": "sha256:77fd4710def9ec3c0f6225800e0235f15a425abd4a8b03559127fcd782612049",
     "LICENSES/protein-ttt/LICENSE": "sha256:bb01e7d5554f9e2e117172e56551452f68a7818df7bc8e71cd7a776a1d4ba3df",
     "LICENSES/protein-ttt/PROVENANCE.md": "sha256:dc641c37353c2efd50ccbdb316ca4aae495ec02c1563e0e15bac92f75fc482e5",
-    "README.md": "sha256:8eb4b722439ea0691e41111c33af8784bae13898595f121e7648ff28bac8a1e1",
+    "README.md": "sha256:e294773edd87bf543956fae98b8ba012a24c9ad79df758beedbd856500cd0c40",
     "THIRD_PARTY_NOTICES.md": "sha256:25704b3c76404696cae52e7fca13088d329f70f412687340351259e86cd62baa",
-    "config.json": "sha256:0b2d1dc1a12593fa81777fd6da4c980017fb3ec18a1060f493fa53f299bcb4e9",
+    "config.json": "sha256:a7311f496acad420a6e5499474b95fa4493f1086ef0372ab7226a958d4377370",
     "fastplms/__init__.py": "sha256:8503e02debf24abc6d33c1913ff3eb9f245e63f7e62017e220d072c7393fdf2c",
     "fastplms/attention/__init__.py": "sha256:ab3c0b6156968f418f3c97ea664959cb978f4137717b4362213bcfe25ab5e3d6",
     "fastplms/attention/_core.py": "sha256:09d0f71a5ea2ad2138c9441f6a2f72bbf8e72729be41d8db737b163f36ce47da",
@@ -46,7 +46,7 @@
     "fastplms/models/esmfold2/esmfold2_parsing.py": "sha256:6ca4a2d85d6158cadcde79bf8d1820c54eeb2160e00b99643c29e270d0a3f870",
     "fastplms/models/esmfold2/esmfold2_predicted_aligned_error.py": "sha256:b76577c9e17bdef5417fea35c6ac134cff2ba949a1c9cb389323134541e52d04",
     "fastplms/models/esmfold2/esmfold2_prepare_input.py": "sha256:f2b5f63a451555d41944e5abc8b85a978edfb456ce493821b313f8343d3eed19",
-    "fastplms/models/esmfold2/esmfold2_processor.py": "sha256:c3be4ef4737f86c7f7969749ad1ad5e483fb6e2a7a2cb7eeef52afad28d3c40e",
+    "fastplms/models/esmfold2/esmfold2_processor.py": "sha256:4eaed444d316bdb31a72b518c64ad11ede6554b88f22a68fd9bf109702529994",
     "fastplms/models/esmfold2/esmfold2_protein_chain.py": "sha256:e87d7cbfab56d0cd420d8ce8b15a23872cdfe6c1036f2467555e57691af172af",
     "fastplms/models/esmfold2/esmfold2_protein_complex.py": "sha256:ee78539192414000b78da7285175f0316d79f549c179678638d009054b4985f7",
     "fastplms/models/esmfold2/esmfold2_protein_structure.py": "sha256:d77254171dd7dc7e269693091e51aa04a7cd8dc4a442d425fab7d922c1b1308d",
@@ -55,28 +55,28 @@
     "fastplms/models/esmfold2/esmfold2_system.py": "sha256:4a0b339bc7446a9ca255f6005ecb8d385c69b7a3e57170c5a48e372bb2340d13",
     "fastplms/models/esmfold2/esmfold2_types.py": "sha256:8963440d40e3143b984679a716a5ac8327fa09692765836dea179c41e2bc9f5d",
     "fastplms/models/esmfold2/esmfold2_utils_types.py": "sha256:b846d61a383fc55091231ec8bc2b829045a4edb52dcd3907c30b81b6057fd4eb",
-    "fastplms/models/esmfold2/modeling_esmfold2.py": "sha256:6b6e4629ad04c561ea8c5aa674eae59f7d4284e97c4d027a5a5ef50db258202c",
-    "fastplms/models/esmfold2/modeling_esmfold2_common.py": "sha256:d9545f880ba6e92b7ed59a7894a4172fb7470bd4b2077fe8ed570ba7cf8bc614",
-    "fastplms/models/esmfold2/modeling_esmfold2_experimental.py": "sha256:22e24d76869045f545ef192e748d5ba995078c8758becae62a284f68da82e718",
+    "fastplms/models/esmfold2/modeling_esmfold2.py": "sha256:0443ef06718c60b01a7a74874f9a33019926820ca512d99cd286eda958259929",
+    "fastplms/models/esmfold2/modeling_esmfold2_common.py": "sha256:f8c4f7349ecd8f2671be9ff2d58c8729d07426c0b8fc7ed047b05a168e50d29d",
+    "fastplms/models/esmfold2/modeling_esmfold2_experimental.py": "sha256:de8d215d137744bf4d2b1c7bdf98617d7c20a425ea25d9c74f88b34d3303ca84",
     "fastplms/models/esmfold2/protein_reference_geometry.json": "sha256:c726d81b928df0005b2d6877e2c0d5ef3ebfb11744e04ca8972388057dc565c0",
     "fastplms/models/esmfold2/protein_utils.py": "sha256:a51cdad1769620389ca58d6782bc787ed70cadb1f310689cf3dcd1d2d2fa92be",
     "fastplms/models/esmfold2/reproducibility.py": "sha256:2027064e3eea918f535d9b60734a3bc9288478f5b957847266d328dcfca88443",
     "fastplms/models/ttt.py": "sha256:f1d7f9298a930f4c40f8e321ffa08b09d1c5b8492bca623bdc7bfd03ec85774d",
     "fastplms/registry.py": "sha256:cb37799dae5a4c0c780089ab2cca4cc56300c1c5f515b275892aa4e1c4b97eb6",
     "fastplms/runtime.py": "sha256:9521c37fcd168bf9a4137277fb73672403ccf30d70d4845aa879fb22f42ddeb6",
-    "fastplms_bundle.py": "sha256:4cf8de3fd1f7126663a22276e4fe16f75282e49089c027cbeb25987831ce6ca1",
-    "modeling_fastplms.py": "sha256:a21e8b519dd202ac28208960221cc0264889425c0c71631003ae9422694ab181",
+    "fastplms_bundle.py": "sha256:b09ae39e6ef318b847e3ff9ca8b5c0dd48a48bd1576994c71c0b5167189e40f2",
+    "modeling_fastplms.py": "sha256:c9a41e763b36487b36afb7ec4836a1dd60c75bd12b2a06dd020a16a4ebdcc912",
     "requirements.txt": "sha256:accda6ebc17b5bc073fff111d29a25fe8fbe0b66b77687aee51100737ccdd8a8"
   },
   "model_id": "esmfold2_fast",
   "redistributable": true,
-  "release_tool_revision": "e6dd397a9ad368c998d714f6bd64d40b533d1ed1",
-  "release_tool_sha256": "6d335c05aa49a232086a816deb25d248d1490529e5783acc3355b9bc6f03e0c2",
-  "runtime_bundle_sha256": "4646da0fba1e0aa8302ca3e78bbf9f00d1123144940f4f377010af0b21ab86b8",
-  "runtime_revision": "e6dd397a9ad368c998d714f6bd64d40b533d1ed1",
+  "release_tool_revision": "b5669feacd99580cab27c9e969092fe9c647b9c6",
+  "release_tool_sha256": "19f8b87151866ea1bec0274bb99c98890bfbb70887276e57d3990a96d5c6b553",
+  "runtime_bundle_sha256": "6e81dad15aa6df4f53633997daf8747aa49a4afbd030a65937a6f3a7d8159273",
+  "runtime_revision": "b5669feacd99580cab27c9e969092fe9c647b9c6",
   "schema_version": 2,
   "scope": "runtime-only",
-  "source_tree_sha256": "332058883bd49a8bfda235a1d983bef21f36de1987e44c1f05ffd4247d0d67f1",
+  "source_tree_sha256": "db6fe0054a134e77d111381b520917ffd0d795e4be0bc2cd0e97d43c59b5f3ee",
   "weights": {
     "repo_id": "Synthyra/ESMFold2-Fast",
     "revision": "407875bfcaa42552bfcb25acd67ee1888b790170"