faustomilletari · Sep 16, 2026 · Sep 15, 2026 · Sep 15, 2026
diff --git a/README.md b/README.md
@@ -37,15 +37,31 @@ The **[ESM Atlas](https://biohub.ai/esm/protein/atlas)** is a map of 6.8 billion
 For information on using ESM3, see the [ESM3 README](https://github.com/Biohub/esm/blob/main/_assets/ESM3_README.md).
 
 ## Table of Contents
-
+- [ESM on PyPI](#esm-on-pypi)
+- [ESM Package vs. Hugging Face Transformers Usage](#esm-package-vs-hugging-face-transformers-usage)
 - [ESMC](#esmc)
 - [ESMC Sparse Autoencoders](#esmc-sparse-autoencoders)
 - [ESMFold2](#esmfold2)
-- [Batch Inference](#batch-inference)
+- [Parallel Job Execution](#parallel-job-execution)
 - [Frontier-Safety](#frontier-safety)
 - [Licenses](#licenses)
 - [Citations](#citations)
 
+## ESM on PyPI
+<a name="esm-on-pypi"></a>
+
+[`esm v3.4.post1`](https://github.com/Biohub/esm/releases/tag/v3.4.1.post1) is now available on PyPI with full support for ESMC and ESMFold2.
+
+## ESM Package vs. Hugging Face Transformers Usage
+<a name="esm-package-vs-hugging-face-transformers-usage"></a>
+ESMC and ESMFold2 are now available in Hugging Face Transformers starting with [v5.16.0](https://github.com/huggingface/transformers/releases/tag/v5.16.0) (and subsequent releases, including v5.16.1).
+
+The ESM package and the Hugging Face Transformers implementation are suited for differing purposes:
+
+|  | Esm Package | Hugging Face Transformers |
+| :---- | :---- | :---- |
+| Use cases | <ul><li> Use for production inference because it contains the optimized dependencies and inference path. </li><li> Full support for ESMC and ESMFold2 along with [Fold-CP](https://arxiv.org/abs/2603.14806) for folding longer sequences. | <ul><li>Access ESMC and ESMFold2 using the familiar Transformers API alongside thousands of other models, lowering the barrier for experimentation, benchmarking, and integration into existing ML workflows. </li><li>Use ESM models within the broader Hugging Face ecosystem, including downstream libraries and tooling. Quickly evaluate the models Note that the Transformers implementation uses a different dependency stack to conform to Transformers requirements. |
+
 ## ESMC
 <a name="esmc"></a>
 
@@ -174,6 +190,9 @@ output.sae_outputs["layer60"]  # sparse.coo tensor
 print(output.sae_outputs["layer60"].shape)
 
 ```
+
+SAEs come in two kinds. Most take the hidden state `h[N]`; some take the residual update `h[N] - h[N-1]`. A repo declares which via `use_residual_update_instead_of_states` in its `config.json`, and the backbone feeds each attached SAE accordingly.
+
 ### Running SAEs Through The Biohub Platform
 
 For a tutorial on using SAEs using the Biohub Platform, see [here](https://github.com/Biohub/esm/blob/main/cookbook/tutorials/esmc_sae_feature_interpretation.ipynb).
@@ -294,6 +313,7 @@ with open("result.cif", "w") as f:
 For tutorials on how to use ESMFold2, see our [tutorials](https://github.com/Biohub/esm/tree/main/cookbook/tutorials).
 
 ## Parallel Job Execution
+<a name="parallel-job-execution"></a>
 For jobs that require processing multiple inputs, the Parallel Executor sends parallel requests, executing them concurrently and efficiently while respecting rate limits and adapting to request latency. The following example provides an example of using the parallel_executor context manager when embedding sequences.
 
 ```py

diff --git a/cookbook/local/open_generate.ipynb b/cookbook/local/open_generate.ipynb
@@ -374,16 +374,19 @@
     "view = py3Dmol.view(width=500, height=500)\n",
     "view.addModel(helix_shortening_chain.to_pdb_string(), \"pdb\")\n",
     "view.setStyle({\"cartoon\": {\"color\": \"lightgrey\"}})\n",
-    "helix_region = np.arange(38, 111)  # zero-indexed\n",
+    "helix_region = np.arange(47, 120)  # zero-indexed\n",
     "view.addStyle(\n",
     "    {\"resi\": (helix_region + 1).tolist()}, {\"cartoon\": {\"color\": \"lightblue\"}}\n",
     ")\n",
     "view.zoomTo()\n",
     "view.show()\n",
-    "helix_shortening_ss8 = \"CCCSHHHHHHHHHHHTTCHHHHHHHHHHHHHTCSSCCCCHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHTTCHHHHHHHHHHHHHHHHHHHHHHHHHHHHIIIIIGGGCCSHHHHHHHHHHHHHHHHHHHHHCCHHHHHHHHHHHHHHHHHHHHHHHHHSCTTCHHHHHHHHHHHHHIIIIICCHHHHHHHHHHHHHHHHTTCTTCCSSHHHHHHHHHHHHHHHHHHHC\"\n",
+    "helix_shortening_ss8 = \"CCCCCCCCCCCSHHHHHHHHHHHTTCHHHHHHHHHHHHHTCSSCCCCHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHHTTCHHHHHHHHHHHHHHHHHHHHHHHHHHHHIIIIIGGGCCSHHHHHHHHHHHHHHHHHHHHHCCHHHHHHHHHHHHHHHHHHHHHHHHHSCTTCHHHHHHHHHHHHHIIIIICCHHHHHHHHHHHHHHHHTTCTTCCSSHHHHHHHHHHHHHHHHHHHCCCCCCCCCCCCCCCCCC\"\n",
     "print(\n",
     "    \"Secondary structure of protein: (H: Alpha Helix, E: Beta Strand, C: Coil) \\n\\t\",\n",
     "    helix_shortening_ss8,\n",
+    ")\n",
+    "print(\n",
+    "    f\"SS8 length: {len(helix_shortening_ss8)}, Sequence length: {len(helix_shortening_chain.sequence)}\"\n",
     ")"
    ]
   },

diff --git a/cookbook/local/raw_forwards.py b/cookbook/local/raw_forwards.py
@@ -18,7 +18,10 @@
 from esm.utils.types import FunctionAnnotation
 
 
+# The open weights are bf16, so calling `forward` directly needs autocast; the SDK
+# generate/logits paths set it up for you.
 @torch.no_grad()
+@torch.autocast("cuda", dtype=torch.bfloat16)
 def inverse_folding_example():
     tokenizer = EsmSequenceTokenizer()
     encoder = ESM3_structure_encoder_v0("cuda")
@@ -47,6 +50,7 @@ def inverse_folding_example():
 
 
 @torch.no_grad()
+@torch.autocast("cuda", dtype=torch.bfloat16)
 def conditioned_prediction_example():
     tokenizers = get_esm3_model_tokenizers()
 

diff --git a/cookbook/snippets/fold_invfold.py b/cookbook/snippets/fold_invfold.py
@@ -108,7 +108,9 @@ def inverse_fold(
         print(
             "Please export your Forge/Biohub Platform API key as ESM_API_KEY environment variable."
         )
-    client = SequenceStructureForgeInferenceClient(token=os.environ["ESM_API_KEY"])
+    client = SequenceStructureForgeInferenceClient(
+        model="esm3-medium-2024-08", token=os.environ["ESM_API_KEY"]
+    )
     esm3_client = ESM3ForgeInferenceClient(
         model="esm3-medium-2024-08", token=os.environ["ESM_API_KEY"]
     )

diff --git a/cookbook/tutorials/esmc_finetune.ipynb b/cookbook/tutorials/esmc_finetune.ipynb
@@ -327,10 +327,11 @@
     "\n",
     "# layernorm_qkv and ffn fused weights are bare nn.Parameter tensors (not Linear\n",
     "# submodules), so they are reached via target_parameters rather than target_modules.\n",
+    "# peft rejects a non-zero lora_dropout on target_parameters.\n",
     "lora_config = LoraConfig(\n",
     "    r=LORA_RANK,\n",
     "    lora_alpha=LORA_ALPHA,\n",
-    "    lora_dropout=0.01,\n",
+    "    lora_dropout=0.0,\n",
     "    target_modules=[\"out_proj\"],\n",
     "    target_parameters=[\"layernorm_qkv.weight\", \"ffn.fc1_weight\", \"ffn.fc2_weight\"],\n",
     "    modules_to_save=[\"classifier\"],\n",

diff --git a/cookbook/tutorials/esmc_layer_sweep.ipynb b/cookbook/tutorials/esmc_layer_sweep.ipynb
@@ -480,7 +480,7 @@
     "N_SPLITS = 5  # number of cross-validation folds\n",
     "RANDOM_STATE = 44\n",
     "\n",
-    "classes = sorted(pd.unique(labels))\n",
+    "classes = sorted(set(labels))\n",
     "y = np.array([classes.index(l) for l in labels])\n",
     "\n",
     "skf = StratifiedKFold(n_splits=N_SPLITS, shuffle=True, random_state=RANDOM_STATE)\n",

diff --git a/cookbook/tutorials/esmfold2_local_applesilicon.ipynb b/cookbook/tutorials/esmfold2_local_applesilicon.ipynb
@@ -512,9 +512,9 @@
    "source": [
     "## 6. If you want to go further\n",
     "\n",
-    "- **A desktop app.** [iESMFold2](https://github.com/faustomilletari/iESMFold2)\n",
-    "  wraps this in a native macOS app — sequence boxes, an interactive viewer, a\n",
-    "  model picker.\n",
+    "- **A desktop app.** The [ESMFold2 Desktop app for Apple\n",
+    "  Silicon](https://github.com/faustomilletari/iESMFold2) wraps this up with\n",
+    "  sequence boxes, an interactive viewer and a model picker.\n",
     "- **Ask for several answers at once.** `num_diffusion_samples=N` returns N\n",
     "  structures; rank them by confidence and keep the best.\n",
     "- **Think harder.** `num_loops` controls how many refinement passes the model\n",
@@ -523,7 +523,7 @@
     "  most of a 48 GB machine.\n",
     "\n",
     "**Links** — [MLX port](https://github.com/faustomilletari/mlx-lm) ·\n",
-    "[desktop app](https://github.com/faustomilletari/iESMFold2) ·\n",
+    "[ESMFold2 Desktop app](https://github.com/faustomilletari/iESMFold2) ·\n",
     "[esm](https://github.com/Biohub/esm) ·\n",
     "[MLX](https://github.com/ml-explore/mlx)\n"
    ]

diff --git a/cookbook/tutorials/esmprotein.ipynb b/cookbook/tutorials/esmprotein.ipynb
@@ -178,8 +178,10 @@
     "    This uses all Alanines\n",
     "    \"\"\"\n",
     "    protein_with_same_coords = ESMProtein(coordinates=coordinates)\n",
-    "    # pdb with all alanines\n",
-    "    pdb_string = protein_with_same_coords.to_pdb_string()\n",
+    "    # pdb with all alanines; to_pdb_string() needs a sequence, to_protein_chain() does not\n",
+    "    pdb_string = (\n",
+    "        protein_with_same_coords.to_protein_chain().infer_oxygen().to_pdb_string()\n",
+    "    )\n",
     "    return visualize_pdb(pdb_string)\n",
     "\n",
     "\n",

diff --git a/esm/__init__.py b/esm/__init__.py
@@ -1 +1 @@
-__version__ = "3.4.1"
+__version__ = "3.4.1.post1"
diff --git a/esm/models/esmc/kernels.py b/esm/models/esmc/kernels.py
@@ -61,7 +61,7 @@
     FLASH_ATTN_ROTARY_INSTALLED = False
 
 
-if not TE_INSTALLED:
+if torch.cuda.is_available() and not TE_INSTALLED:
     logger.warning(
         "ESMC: Transformer Engine is not installed; falling back to "
         "pure-PyTorch LayerNorm+Linear / LayerNorm+MLP. Outputs will differ "
@@ -73,7 +73,7 @@
         "reduction LayerNorm."
     )
 
-if not XFORMERS_INSTALLED and not FLASH_ATTN_INSTALLED:
+if torch.cuda.is_available() and not (XFORMERS_INSTALLED or FLASH_ATTN_INSTALLED):
     logger.warning(
         "ESMC: neither xformers nor flash-attn is installed; falling back "
         "to PyTorch ``F.scaled_dot_product_attention``. The attention "

diff --git a/esm/models/esmc/model.py b/esm/models/esmc/model.py
@@ -332,6 +332,29 @@ def add_sae_models(self, sae_models: list[EsmcSaeLayer]) -> None:
                 )
             self._sae_models[key] = layer
 
+    def _sae_layers_to_collect(self) -> list[int]:
+        """Backbone layers the registered SAEs need collected.
+
+        A residual-update SAE at layer ``N`` needs ``N - 1`` too, since
+        ``h[N] - h[N-1]`` is differenced after collection.
+
+        Returns
+        -------
+        list[int]
+            Sorted layer indices.
+        """
+        needed: set[int] = set()
+        for model_name, sae_module in self._sae_models.items():
+            assert isinstance(sae_module, EsmcSaeLayer)
+            layer_num = self._get_sae_layer_num_requested(model_name)
+            needed.add(layer_num)
+            if (
+                sae_module.params.use_residual_update_instead_of_states
+                and layer_num > 0
+            ):
+                needed.add(layer_num - 1)
+        return sorted(needed)
+
     def _get_sae_layer_num_requested(self, model_name: str) -> int:
         match = self._SAE_KEY_RE.fullmatch(model_name)
         assert match is not None, (
@@ -363,6 +386,21 @@ def _get_sae_outputs(
             layer_idx = layer_to_idx[requested_layer]
             layer_states = hidden_states[layer_idx].clone().to(self.device)
 
+            # Layer 0 has no predecessor and stays the raw state, as in training.
+            if (
+                layer.params.use_residual_update_instead_of_states
+                and requested_layer > 0
+            ):
+                predecessor = requested_layer - 1
+                if predecessor not in layer_to_idx:
+                    raise KeyError(
+                        f"SAE {model_name!r} uses residual updates and needs "
+                        f"layer {predecessor} collected alongside layer "
+                        f"{requested_layer}, but the collected layers are "
+                        f"{layers_to_collect}."
+                    )
+                layer_states -= hidden_states[layer_to_idx[predecessor]].to(self.device)
+
             sae_out = layer.get_sae_output(layer_states, token_mask)
             features = sae_out.feature_magnitudes.detach()
 
@@ -429,9 +467,7 @@ def forward(
                 range(self.config.num_hidden_layers + 1)
             )
         elif output_sae:
-            layers_to_collect = sorted(
-                {self._get_sae_layer_num_requested(name) for name in self._sae_models}
-            )
+            layers_to_collect = self._sae_layers_to_collect()
         else:
             layers_to_collect = []
 

diff --git a/esm/models/esmc/sae.py b/esm/models/esmc/sae.py
@@ -32,6 +32,7 @@ class EsmcSaeParams:
     codebook_dim: int = 65536
     k: int = 64
     layer: int = 0
+    use_residual_update_instead_of_states: bool = False
 
 
 @dataclass
@@ -55,6 +56,10 @@ class EsmcSaeConfig:
         Top-k sparsity per SAE.
     available_layers : list[int]
         Which backbone-layer indices the repo ships.
+    use_residual_update_instead_of_states : bool
+        ``True`` when these SAEs were trained on the residual update
+        ``h[i] - h[i-1]`` rather than the hidden state ``h[i]``. Mirrors the
+        training-time flag of the same meaning.
     """
 
     model_type = "esmc_sae"
@@ -63,6 +68,7 @@ class EsmcSaeConfig:
     codebook_dim: int = 65536
     k: int = 64
     available_layers: list[int] | None = None
+    use_residual_update_instead_of_states: bool = False
 
     def __post_init__(self) -> None:
         self.available_layers = (
@@ -78,6 +84,10 @@ def from_pretrained(cls, directory: str | os.PathLike) -> "EsmcSaeConfig":
             codebook_dim=raw["codebook_dim"],
             k=raw["k"],
             available_layers=raw.get("available_layers"),
+            # The 97 repos published before this flag existed carry no key.
+            use_residual_update_instead_of_states=raw.get(
+                "use_residual_update_instead_of_states", False
+            ),
         )
 
     def save_pretrained(self, directory: str | os.PathLike) -> None:
@@ -285,6 +295,9 @@ def initialize_layers(
                 codebook_dim=self.config.codebook_dim,
                 k=self.config.k,
                 layer=layer_idx,
+                use_residual_update_instead_of_states=(
+                    self.config.use_residual_update_instead_of_states
+                ),
             )
             # Build on the meta device so we don't allocate weights that
             # ``load_state_dict`` would immediately overwrite.

diff --git a/esm/models/esmfold2/experimental.py b/esm/models/esmfold2/experimental.py
@@ -34,6 +34,7 @@
 from esm.models.esmc import EsmcModel
 from esm.models.esmc.checkpoint_layout import published_to_native_subtree
 from esm.models.esmfold2.config import EsmFold2Config
+from esm.models.esmfold2.constants import MOL_TYPE_NONPOLYMER
 from esm.models.esmfold2.layers import (
     CHAR_VOCAB_SIZE,
     MAX_ATOMIC_NUMBER,
@@ -63,7 +64,6 @@
 from esm.models.hub import HubPreTrainedModel, resolve_model_dir
 
 _EPS = 1e-5
-_NONPOLYMER_ID: int = 3
 
 
 # ===========================================================================
@@ -246,7 +246,7 @@ def forward(
         # Complex ipLDDT (interface-weighted)
         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()
+        is_ligand = (expanded_type == MOL_TYPE_NONPOLYMER).float()
         inter_chain = (
             expanded_asym.unsqueeze(-1) != expanded_asym.unsqueeze(-2)
         ).float()
@@ -892,6 +892,8 @@ def forward(
         max_inference_sigma: float | None = 256.0,
         seed: int | None = None,
         provide_soft_sequence_to_msa_and_profile: bool = True,
+        disto_cond: Tensor | None = None,
+        disto_cond_mask: Tensor | None = None,
         include_embeddings: bool = False,
         **unused_features: Tensor,
     ) -> dict[str, Tensor]:
@@ -909,6 +911,15 @@ def forward(
                 f"argument(s) {unexpected}."
             )
 
+        # Named rather than left to **kwargs so conditioning cannot be swallowed
+        # and the fold silently run unconditioned. This model has no
+        # disto_conditioning_proj that would consume these.
+        if disto_cond_mask is not None and disto_cond_mask.any():
+            raise NotImplementedError(
+                "distogram conditioning is not implemented for ESMFold2; "
+                "fold without it."
+            )
+
         tok_mask = token_attention_mask
         atm_mask = atom_attention_mask
         disto_idx = distogram_atom_idx

diff --git a/esm/models/esmfold2/model.py b/esm/models/esmfold2/model.py
@@ -34,6 +34,7 @@
 from esm.models.esmc import EsmcModel
 from esm.models.esmc.checkpoint_layout import published_to_native_subtree
 from esm.models.esmfold2.config import EsmFold2Config, default_module_flags
+from esm.models.esmfold2.constants import MOL_TYPE_NONPOLYMER
 from esm.models.esmfold2.hf_checkpoint import hf_state_dict_to_native, is_hf_layout
 from esm.models.esmfold2.layers import (
     CHAR_VOCAB_SIZE,
@@ -60,7 +61,6 @@
 from esm.models.hub import HubPreTrainedModel, resolve_model_dir
 
 _EPS = 1e-6
-_NONPOLYMER_ID = 4
 
 # Default for the triangle / OPM / pair-transition L² ops. Caps peak memory
 # so L≈2k folds on an 80 GB GPU (~76 GB peak at chunk=128 for L=1438;
@@ -305,7 +305,7 @@ def _finish(
 
         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()
+        is_ligand = (expanded_type == MOL_TYPE_NONPOLYMER).float()
         inter_chain = (
             expanded_asym.unsqueeze(-1) != expanded_asym.unsqueeze(-2)
         ).float()