smaller

README.md

@@ -8,8 +8,7 @@ python_version: "3.12"
 app_file: app.py
 pinned: false
 license: mit
-short_description: Protein structure prediction with ESMFold2 (ZeroGPU)
-suggested_hardware: zero-a10g
+short_description: Protein structure prediction with ESMFold2-300 (ZeroGPU)
 tags:
   - biology
   - protein
@@ -21,39 +20,59 @@ tags:
 # Foldy
 
 Predicts 3D protein structure from amino acid sequences using
-[ESMFold2](https://huggingface.co/biohub/ESMFold2-Fast) (Biohub), a 7B-parameter
-diffusion-based structure predictor. Runs on Hugging Face's ZeroGPU.
+[Synthyra/ESMFold2-300](https://huggingface.co/Synthyra/ESMFold2-300), a 300M-backbone
+experimental checkpoint from Biohub. Runs on Hugging Face's ZeroGPU.
+
+## ⚠️ Third-party remote code — read before touching `MODEL_NAME`
+
+This app loads its model with `trust_remote_code=True`, from `Synthyra/ESMFold2-300` —
+**not** an official Biohub repo. That model's `modeling_fastplms.py` doesn't contain plain
+readable model code: it base85-decodes a bundled ~400KB payload (`fastplms_bundle.py`) into
+a temp directory and imports it at runtime. That's a materially bigger trust step than a
+typical `trust_remote_code` model with inspectable source — the actual code only becomes
+readable after it's already unpacked and about to run. This was a deliberate, explicit
+decision (see conversation history), not a default anyone should assume is fine to repeat
+elsewhere without the same consideration.
+
+Why we ended up here: the *official* `biohub/ESMFold2-Fast` (7B params) reliably OOMs on
+large complexes even on ZeroGPU's `xlarge` (96GB) tier, because of the model's O(L²) pairwise
+representation. `Synthyra/ESMFold2-300` wraps a genuinely smaller (~300M) Biohub checkpoint
+with a much lower memory footprint, at the cost of: third-party trust (above), no confidence
+head (no pLDDT/pTM/ipTM), and "experimental" status per Biohub's own labeling (thin
+validation — "passed Docker BF16 reference comparison on one compact Protein G sequence").
 
 ## Input
 
-A single raw amino acid sequence, or a single-record FASTA (the header is ignored).
-Length is capped (see `MAX_RESIDUES` in [app.py](app.py)) to keep inference within the
-ZeroGPU time budget. Multi-chain/complex folding isn't wired up here — it needs the
-`esm` PyPI package's `StructurePredictionInput` builder, which pulls in a dependency
-tree (a hard `torch<2.12` pin, `cuequivariance-torch`, ...) that conflicts with the
-versions this Space's ZeroGPU/Gradio build forces. This app instead uses ESMFold2 as
-natively vendored inside `transformers` (`transformers.models.esmfold2`), which only
-supports single-chain protein folding but has a clean, conflict-free dependency set.
+Up to 6 protein chains (fixed rows in the UI), each with a copies count — set copies > 1 for
+a homooligomer. No cap on total residues; oversized inputs just hit the
+`@spaces.GPU(duration=300)` timeout instead of a friendly error, by design (see conversation
+history — no evidence-based cap has been set yet).
 
 ## Output
 
-- An interactive 3D viewer (py3Dmol), colored by predicted confidence (pLDDT)
-- Mean pLDDT and pTM
-- A downloadable PDB file
+- An interactive 3D viewer (py3Dmol), colored by chain (no confidence scores to color by)
+- A downloadable mmCIF file
+- No pLDDT/pTM/ipTM — this checkpoint's confidence head is disabled
 
 ## Notes
 
-- Uses the `ESMFold2-Fast` checkpoint (single-sequence, no MSA) for speed. Swap
-  `MODEL_NAME` in [app.py](app.py) for `biohub/ESMFold2` if you want MSA-conditioned
-  predictions (slower, and you'll need to supply an MSA yourself — the native
-  `transformers` path here doesn't build one for you).
-- Model weights load once at startup in bfloat16 and are moved to the GPU inside the
-  `@spaces.GPU`-decorated function, per the ZeroGPU pattern.
-- "Recycling loops" and "diffusion sampling steps" trade accuracy for speed — raise
-  them for harder targets, lower them for a faster demo.
-- Do **not** add `esm`, `torch`, `spaces`, or `gradio` to `requirements.txt`: the
-  platform force-pins `torch`/`spaces`/`gradio` during the ZeroGPU build, and `esm`'s
-  own pins (`torch<2.12`, `transformers<5.0`) directly conflict with them.
+- `model.set_chunk_size(32)` and `esmc_precision="bf16"` are both part of this checkpoint's
+  own documented usage (not something we bolted on) — they trade some accuracy/speed for
+  lower memory.
+- Folding goes through `model.input_types.StructurePredictionInput` /
+  `.ProteinInput` and `model.fold(...)` / `model.result_to_cif(...)` — all native to this
+  checkpoint's bundle, not the hand-rolled multi-chain featurization used in earlier
+  iterations of this app.
+- GPU size is the ZeroGPU default (`large`, 48GB) — the 300M model's weights are ~685MB, so
+  this should have real headroom versus the old 7B setup. Bump `@spaces.GPU(size="xlarge")`
+  if you still hit OOM on large complexes.
+- "Recycling loops" and "diffusion sampling steps" trade accuracy for speed — raise them for
+  harder targets, lower them for a faster demo. Defaults (3 loops / 15 steps) match this
+  checkpoint's own documented config, not the smaller/faster defaults used for the old model.
+- Do **not** add `spaces` or `gradio` to `requirements.txt` — the platform force-pins those
+  during the ZeroGPU build; adding them again risks a resolver conflict. `torch`/`transformers`
+  *are* pinned here deliberately (matching what `Synthyra/ESMFold2-300` requires), verified to
+  resolve cleanly against the platform's forced pins before being added.
 
 ## Local development
 
@@ -62,5 +81,6 @@ pip install -r requirements.txt
 python app.py
 ```
 
-Folding requires a CUDA GPU; on Spaces, ZeroGPU provides this on demand via the
-`spaces` package. Locally, you need your own GPU for `model.to("cuda")` to work.
+Folding requires a CUDA GPU; on Spaces, ZeroGPU provides this on demand via the `spaces`
+package. Locally, you need your own GPU — `device_map="cuda"` in `app.py` assumes one is
+present at import time.

app.py

@@ -1,4 +1,9 @@
-"""Foldy — protein structure prediction with ESMFold2, on HF ZeroGPU."""
+"""Foldy — protein structure prediction with ESMFold2-300, on HF ZeroGPU.
+
+Uses the Synthyra/ESMFold2-300 FastPLMs bundle (trust_remote_code=True), which wraps
+Biohub's experimental 300M-backbone ESMFold2 checkpoint. This is third-party remote code —
+see the README for what that means before touching MODEL_NAME.
+"""
 
 import os
 import re
@@ -9,13 +14,9 @@ import gradio as gr
 import py3Dmol
 import spaces
 import torch
-from transformers.models.esmfold2.modeling_esmfold2 import EsmFold2Model
-from transformers.models.esmfold2.protein_utils import output_to_pdb, prepare_protein_features
+from transformers import AutoModel
 
-# The -Fast checkpoint is single-sequence (no MSA), which keeps inference time
-# reasonable on a shared ZeroGPU slot. Swap for "biohub/ESMFold2" for MSA-conditioned
-# accuracy on hard targets (slower, and needs an MSA pipeline this app doesn't provide).
-MODEL_NAME = "biohub/ESMFold2-Fast"
+MODEL_NAME = "Synthyra/ESMFold2-300"
 NUM_ROWS = 6
 VALID_AMINO_ACIDS = set("ACDEFGHIKLMNPQRSTVWYXBZUO")
 
@@ -23,11 +24,16 @@ EXAMPLE_UBIQUITIN = "MQIFVKTLTGKTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDY
 EXAMPLE_CAP = "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG"
 
 print(f"Loading {MODEL_NAME} (this happens once at startup)...")
-model = EsmFold2Model.from_pretrained(MODEL_NAME).eval()
-# Per HF's ZeroGPU docs, models must move to "cuda" at module load time (not inside an
-# @spaces.GPU function) — a CUDA emulation mode makes this safe even without a real GPU
-# attached yet, and it's more efficient than a per-call transfer.
-model.to("cuda")
+model = AutoModel.from_pretrained(
+    MODEL_NAME,
+    trust_remote_code=True,
+    dtype=torch.float32,
+    device_map="cuda",
+    esmc_precision="bf16",
+    attn_implementation="sdpa",
+).eval()
+model.set_chunk_size(32)
+input_types = model.input_types
 print("Model loaded.")
 
 
@@ -38,148 +44,83 @@ def clean_sequence(raw_sequence: str) -> str:
     return re.sub(r"[^A-Z]", "", sequence)
 
 
-def render_structure(pdb_str: str) -> str:
+def render_structure(cif_str: str) -> str:
     view = py3Dmol.view(width=600, height=500)
-    view.addModel(pdb_str, "pdb")
-    view.setStyle({"cartoon": {"colorscheme": {"prop": "b", "gradient": "roygb", "min": 50, "max": 90}}})
+    view.addModel(cif_str, "cif")
+    # No per-residue confidence on this checkpoint (confidence head is disabled), so color by
+    # chain instead of pLDDT.
+    view.setStyle({"cartoon": {"colorscheme": "chain"}})
     view.zoomTo()
     raw_html = view._make_html()
     escaped = raw_html.replace('"', """)
     return f'<iframe style="width: 100%; height: 500px; border: none;" srcdoc="{escaped}"></iframe>'
 
 
-def prepare_multi_chain_protein_features(sequences: list[str], device) -> dict[str, torch.Tensor]:
-    """Combine several single-chain protein featurizations into one complex.
-
-    Reuses ``prepare_protein_features`` per chain (already correct for one chain), then
-    reassigns the complex-level bookkeeping fields per the convention ``EsmFold2Model.forward``
-    documents: ``token_index``/``atom_to_token``/``ref_space_uid`` become complex-global,
-    ``residue_index`` stays chain-local, ``asym_id`` is unique per chain instance, ``entity_id``
-    is shared by chains with an identical sequence, and ``sym_id`` numbers the copies within an
-    entity (so copies of one sequence fold as a homooligomer, not independent monomers).
-    """
-    per_chain = [prepare_protein_features(seq, device="cpu") for seq in sequences]
-
-    entity_id_by_sequence: dict[str, int] = {}
-    next_entity_id = 1
-    copies_seen: dict[int, int] = {}
-    token_offset = 0
-    atom_offset = 0
-
-    for asym_idx, (seq, feats) in enumerate(zip(sequences, per_chain)):
-        entity_id = entity_id_by_sequence.setdefault(seq, next_entity_id)
-        if entity_id == next_entity_id:
-            next_entity_id += 1
-        sym_id = copies_seen.get(entity_id, 0)
-        copies_seen[entity_id] = sym_id + 1
-
-        num_tokens = feats["token_index"].shape[1]
-
-        feats["token_index"] = feats["token_index"] + token_offset
-        feats["asym_id"] = torch.full_like(feats["asym_id"], asym_idx)
-        feats["sym_id"] = torch.full_like(feats["sym_id"], sym_id)
-        feats["entity_id"] = torch.full_like(feats["entity_id"], entity_id)
-        feats["atom_to_token"] = feats["atom_to_token"] + token_offset
-        feats["ref_space_uid"] = feats["ref_space_uid"] + token_offset
-        feats["distogram_atom_idx"] = feats["distogram_atom_idx"] + atom_offset
-
-        token_offset += num_tokens
-        atom_offset += feats["ref_pos"].shape[1]
-
-    total_tokens = token_offset
-
-    def cat(key: str, dim: int = 1) -> torch.Tensor:
-        return torch.cat([f[key] for f in per_chain], dim=dim)
-
-    combined = {
-        "token_index": cat("token_index"),
-        "residue_index": cat("residue_index"),
-        "asym_id": cat("asym_id"),
-        "sym_id": cat("sym_id"),
-        "entity_id": cat("entity_id"),
-        "mol_type": cat("mol_type"),
-        "res_type": cat("res_type"),
-        "input_ids": cat("input_ids"),
-        "attention_mask": cat("attention_mask"),
-        "distogram_atom_idx": cat("distogram_atom_idx"),
-        "ref_pos": cat("ref_pos"),
-        "ref_element": cat("ref_element"),
-        "ref_charge": cat("ref_charge"),
-        "ref_atom_name_chars": cat("ref_atom_name_chars"),
-        "ref_space_uid": cat("ref_space_uid"),
-        "atom_attention_mask": cat("atom_attention_mask"),
-        "atom_to_token": cat("atom_to_token"),
-        "msa": cat("msa", dim=2),
-        "msa_attention_mask": cat("msa_attention_mask", dim=2),
-        "has_deletion": cat("has_deletion", dim=2),
-        "deletion_value": cat("deletion_value", dim=2),
-        "deletion_mean": cat("deletion_mean"),
-        "token_bonds": torch.zeros(1, total_tokens, total_tokens, 1, dtype=torch.float32),
-    }
-    return {k: v.to(device) for k, v in combined.items()}
-
-
-@spaces.GPU(duration=300, size="xlarge")
+@spaces.GPU(duration=300)
 def fold_rows(*args, progress=gr.Progress()):
     *seqs, num_loops, num_sampling_steps = args
 
-    chain_sequences: list[str] = []
-    for raw_seq, copies in zip(seqs[:NUM_ROWS], seqs[NUM_ROWS:]):
+    chains: list[str] = []
+    for row_idx, (raw_seq, copies) in enumerate(zip(seqs[:NUM_ROWS], seqs[NUM_ROWS:])):
         sequence = clean_sequence(raw_seq or "")
         if not sequence:
             continue
         invalid_chars = set(sequence) - VALID_AMINO_ACIDS
         if invalid_chars:
             raise gr.Error(f"Invalid characters in sequence: {', '.join(sorted(invalid_chars))}")
-        chain_sequences.extend([sequence] * max(1, int(copies or 1)))
+        chains.extend([sequence] * max(1, int(copies or 1)))
 
-    if not chain_sequences:
+    if not chains:
         raise gr.Error("Please enter at least one protein sequence.")
 
-    total_residues = sum(len(s) for s in chain_sequences)
+    total_residues = sum(len(s) for s in chains)
 
-    progress(0.2, desc=f"Folding {len(chain_sequences)} chain(s), {total_residues} residues total...")
-    features = prepare_multi_chain_protein_features(chain_sequences, device="cuda")
+    progress(0.2, desc=f"Folding {len(chains)} chain(s), {total_residues} residues total...")
+    chain_ids = [chr(ord("A") + i) for i in range(len(chains))]
+    complex_input = input_types.StructurePredictionInput(
+        sequences=[
+            input_types.ProteinInput(id=chain_id, sequence=seq) for chain_id, seq in zip(chain_ids, chains)
+        ]
+    )
 
     start = time.time()
-    output = model.fold(
-        **features,
-        num_loops=num_loops,
-        num_sampling_steps=num_sampling_steps,
-        num_diffusion_samples=1,
-    )
+    with torch.inference_mode():
+        result = model.fold(
+            complex_input,
+            num_loops=num_loops,
+            num_sampling_steps=num_sampling_steps,
+            num_diffusion_samples=1,
+            seed=0,
+            verbose=False,
+        )
     elapsed = time.time() - start
 
     progress(0.9, desc="Building structure file...")
-    pdb_str = output_to_pdb(output, features)
+    cif_str = model.result_to_cif(result)
 
-    pdb_path = os.path.join(tempfile.gettempdir(), f"foldy_{int(time.time())}.pdb")
-    with open(pdb_path, "w") as f:
-        f.write(pdb_str)
+    cif_path = os.path.join(tempfile.gettempdir(), f"foldy_{int(time.time())}.cif")
+    with open(cif_path, "w") as f:
+        f.write(cif_str)
 
-    viewer_html = render_structure(pdb_str)
-    mean_plddt = float(output["plddt"].mean())
-    ptm = float(output["ptm"].mean())
+    viewer_html = render_structure(cif_str)
     summary_lines = [
-        f"**Chains:** {len(chain_sequences)}",
+        f"**Chains:** {len(chains)}",
         f"**Total length:** {total_residues} residues",
-        f"**Mean pLDDT:** {mean_plddt:.1f} / 100",
-        f"**pTM:** {ptm:.3f}",
+        "**Confidence scores:** unavailable (this checkpoint's confidence head is disabled)",
+        f"**Inference time:** {elapsed:.1f}s",
     ]
-    iptm_tensor = output.get("iptm") if hasattr(output, "get") else None
-    if len(chain_sequences) > 1 and iptm_tensor is not None:
-        summary_lines.append(f"**ipTM:** {float(iptm_tensor.mean()):.3f}")
-    summary_lines.append(f"**Inference time:** {elapsed:.1f}s")
 
-    return viewer_html, "  \n".join(summary_lines), pdb_path
+    return viewer_html, "  \n".join(summary_lines), cif_path
 
 
 with gr.Blocks(title="Foldy — ESMFold2 Protein Structure Prediction") as demo:
     gr.Markdown(
         """
         # 🧬 Foldy
-        Predict 3D protein structure using [ESMFold2](https://huggingface.co/biohub/ESMFold2-Fast)
-        (Biohub). Fill in one row per chain, and set copies > 1 for a homooligomer.
+        Predict 3D protein structure using [ESMFold2-300](https://huggingface.co/Synthyra/ESMFold2-300)
+        (a smaller, 300M-backbone experimental checkpoint from Biohub, via Synthyra's runtime).
+        Fill in one row per chain, and set copies > 1 for a homooligomer. This checkpoint has no
+        confidence head, so pLDDT/pTM aren't available.
         """
     )
 
@@ -206,13 +147,13 @@ with gr.Blocks(title="Foldy — ESMFold2 Protein Structure Prediction") as demo:
                 copies_boxes.append(copies)
 
             with gr.Accordion("Advanced settings", open=False):
-                num_loops = gr.Slider(1, 3, value=1, step=1, label="Recycling loops")
-                num_sampling_steps = gr.Slider(10, 50, value=20, step=5, label="Diffusion sampling steps")
+                num_loops = gr.Slider(1, 3, value=3, step=1, label="Recycling loops")
+                num_sampling_steps = gr.Slider(10, 50, value=15, step=5, label="Diffusion sampling steps")
                 gr.Markdown("Higher values improve accuracy but increase inference time.")
 
             fold_btn = gr.Button("Fold", variant="primary")
             summary_output = gr.Markdown()
-            pdb_output = gr.File(label="Download PDB", interactive=False)
+            cif_output = gr.File(label="Download mmCIF", interactive=False)
 
         with gr.Column(scale=2):
             structure_output = gr.HTML(label="Predicted structure")
@@ -223,7 +164,7 @@ with gr.Blocks(title="Foldy — ESMFold2 Protein Structure Prediction") as demo:
     fold_btn.click(
         fn=fold_rows,
         inputs=[*seq_boxes, *copies_boxes, num_loops, num_sampling_steps],
-        outputs=[structure_output, summary_output, pdb_output],
+        outputs=[structure_output, summary_output, cif_output],
     )
 
 if __name__ == "__main__":

requirements.txt

@@ -1,2 +1,19 @@
-transformers>=5.16
+torch>=2.13,<2.14
+transformers>=5.13,<5.14
+huggingface-hub>=0.34,<2
+tokenizers>=0.22,<0.23
+safetensors>=0.5,<1
+numpy>=1.26,<3
+einops>=0.8,<1
+tqdm>=4.67,<5
+accelerate>=1.10,<2
+biopython>=1.85,<2
+biotite>=1.4,<2
+brotli>=1.1,<2
+msgpack>=1.1,<2
+msgpack-numpy>=0.4.8,<1
+omegaconf>=2.3,<3
+rdkit>=2025.9,<2027
+scipy>=1.15,<2
+zstandard>=0.23,<1
 py3Dmol