esmfold2

README.md

@@ -8,7 +8,7 @@ python_version: "3.12"
 app_file: app.py
 pinned: false
 license: mit
-short_description: Protein structure prediction with ESMFold2-300 (ZeroGPU)
+short_description: Protein structure prediction with ESMFold2 (ZeroGPU)
 tags:
   - biology
   - protein
@@ -20,92 +20,67 @@ tags:
 # Foldy
 
 Predicts 3D protein structure from amino acid sequences using
-[Synthyra/ESMFold2-300](https://huggingface.co/Synthyra/ESMFold2-300), a 300M-backbone
-experimental checkpoint from Biohub. Runs on Hugging Face's ZeroGPU.
+[biohub/ESMFold2-Fast](https://huggingface.co/biohub/ESMFold2-Fast) — the official,
+7B-parameter checkpoint from Biohub, loaded via `transformers`' own native ESMFold2 support
+(`transformers.models.esmfold2`, no `trust_remote_code`, no third-party wrapper). Runs on
+Hugging Face's ZeroGPU.
 
-## ⚠️ Third-party remote code — read before touching `MODEL_NAME`
+## Input
 
-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.
+Up to 6 protein chains (fixed rows in the UI), each with a copies slider (1–16) — set
+copies > 1 for a homooligomer. Chains are combined into one complex via a hand-rolled
+multi-chain featurizer (`prepare_multi_chain_protein_features` in [app.py](app.py)), since the
+native `transformers` path only ships a single-chain helper
+(`prepare_protein_features`) — ours reassigns `asym_id`/`entity_id`/`sym_id` per the
+convention `EsmFold2Model.forward` documents, verified against the model's own
+multi-chain grouping logic (`_compute_lm_hidden_states`). No cap on total residues; oversized
+inputs just hit the `@spaces.GPU(duration=300, size="xlarge")` timeout/OOM instead of a
+friendly error.
 
-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").
+DNA/RNA/ligands aren't supported: that 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 (`gradio` needs `huggingface-hub>=1.16`, `esm`'s own `transformers<5.0` pin needs
+`huggingface-hub<1.0` — genuinely unsatisfiable together). Protein-only, via the version of
+ESMFold2 natively vendored inside `transformers`, is the tradeoff that keeps this Space's
+dependencies conflict-free.
 
-## Input
+## Output
 
-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).
+- An interactive 3D viewer (py3Dmol), colored by predicted confidence (pLDDT)
+- Mean pLDDT, pTM, and (for multi-chain folds) ipTM
+- A downloadable PDB file
 
-## Output
+## Known limitation: large inputs OOM
 
-- 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
+This checkpoint has a dense O(L²) pairwise representation (like AlphaFold/Evoformer-style
+architectures) — memory scales quadratically with total residue count. In practice:
+- Small-to-moderate proteins/complexes (up to several hundred residues) fold reliably.
+- Multi-thousand-residue complexes have reliably OOM'd in testing, even on ZeroGPU's `xlarge`
+  (96GB) tier — this is a known, currently-unaddressed limitation, not a bug to casually "fix."
+  A real fix would mean auditing every large tensor operation in `transformers`' ESMFold2
+  implementation for complete chunking coverage (see the `sfold`/Synthyra-checkpoint chapter of
+  this project's history for how deep that rabbit hole goes on a *different*, less mature
+  checkpoint) — genuinely out of scope for casual iteration.
+- For context: even production AlphaFold3 needs 80GB+ VRAM for large multimers and is
+  described as needing H100/H200-class hardware above ~1500 residues. 10,000-residue multimer
+  folding on a single GPU is not something published tooling demonstrates without novel,
+  paper-worthy techniques (e.g. activation quantization) — not a realistic goal for this app.
 
 ## Notes
 
-- `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.
-- **`set_chunk_size` does not prevent OOM on large complexes by itself.** Confirmed by
-  decoding the model's actual bundled source: the triangle-multiplicative-update's initial
-  linear projection (`proj_bundle`, 4× the pair width) is computed over the full, unchunked
-  `(L, L, ...)` grid *before* chunking is applied to the later einsum step. At L=1888 this one
-  tensor needs ~7.3GB, and it OOM'd even on ZeroGPU's `large` (48GB) tier. We now also call
-  `model.set_kernel_backend("cuequivariance")` at load time, which routes triangle updates
-  through a fused kernel that avoids materializing that tensor at all — but this requires the
-  `cuequivariance-torch`/`cuequivariance-ops-torch-cu13` packages (added to requirements.txt)
-  *and* a CUDA 13 runtime, which we could not confirm ZeroGPU's RTX Pro 6000 Blackwell nodes
-  actually have without testing live. It's wrapped in a try/except that falls back to the
-  chunked (OOM-prone-on-large-inputs) path if unavailable — check the Space logs at startup
-  for "Enabled cuequivariance..." vs "cuequivariance backend unavailable..." to see which path
-  is actually active.
-- **The cuequivariance fix worked, but moved the OOM further down the pipeline, not away.**
-  A second 1888-residue test (after adding cuequivariance) crashed later, in
-  `pair_loop_proj` (a `LayerNorm -> Linear` run once per recycling loop over the full
-  `(L, L, d_pair)` pair tensor) — this op has *no* chunking in the checkpoint's own code at
-  all, unlike `pair_transition`, which we confirmed chunks correctly (the whole computation
-  happens inside the chunk loop, not just part of it — this checkpoint's memory engineering is
-  genuinely inconsistent across ops, not fundamentally incapable of scale; see conversation
-  history for why we now believe production AlphaFold/OpenFold can do 3k+ residues on 48GB but
-  this experimental checkpoint currently can't). We patch `model.pair_loop_proj.forward` at
-  load time with our own chunked version — verified mathematically exact (bitwise-identical
-  output) against the unchunked version in isolation, since LayerNorm-over-last-dim and Linear
-  are both purely position-wise with no cross-`(i,j)` interaction.
-- **Open question, not yet resolved:** at the point `pair_loop_proj` failed, only ~0.4GB was
-  free out of 48GB — meaning cumulative memory from everything *before* it (ESMC backbone,
-  weights, prior recycling iterations) was already near the ceiling. The `pair_loop_proj` fix
-  reduces its own transient peak but may not be enough on its own if upstream accumulation is
-  the real driver. Only a live test will tell; if it still OOMs, the next places to look are
-  whether intermediate tensors from earlier recycling loop iterations are being retained
-  longer than necessary, and whether `num_loops`/`chunk_size` need to go lower for very large
-  inputs.
-- 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.
+- Model weights load once at startup in **float32** (not bfloat16) and move to `cuda` at
+  **module load time** (not inside the `@spaces.GPU` function) — both per HF's current ZeroGPU
+  guidance and hard-won lessons: blanket bf16 casting crashes (`expected Float but found
+  BFloat16`) because the diffusion sampling loop hardcodes fp32 tensors internally regardless
+  of the model's overall dtype.
 - "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.
+  harder targets, lower them for a faster demo.
+- GPU size is `xlarge` (96GB, 2× ZeroGPU quota cost) — bumped up from the `large` default after
+  the 7B model's fp32 weights (~28GB) left too little headroom for anything but small folds.
+- 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.
 
 ## Local development
 
@@ -115,5 +90,5 @@ 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 — `device_map="cuda"` in `app.py` assumes one is
+package. Locally, you need your own GPU — `model.to("cuda")` at module scope assumes one is
 present at import time.

app.py

@@ -1,9 +1,4 @@
-"""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.
-"""
+"""Foldy — protein structure prediction with ESMFold2, on HF ZeroGPU."""
 
 import os
 import re
@@ -14,9 +9,13 @@ import gradio as gr
 import py3Dmol
 import spaces
 import torch
-from transformers import AutoModel
+from transformers.models.esmfold2.modeling_esmfold2 import EsmFold2Model
+from transformers.models.esmfold2.protein_utils import output_to_pdb, prepare_protein_features
 
-MODEL_NAME = "Synthyra/ESMFold2-300"
+# 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"
 NUM_ROWS = 6
 VALID_AMINO_ACIDS = set("ACDEFGHIKLMNPQRSTVWYXBZUO")
 
@@ -24,48 +23,11 @@ EXAMPLE_UBIQUITIN = "MQIFVKTLTGKTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDY
 EXAMPLE_CAP = "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG"
 
 print(f"Loading {MODEL_NAME} (this happens once at startup)...")
-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)
-try:
-    # set_chunk_size only chunks the triangle-update einsum step — the linear projection
-    # feeding it (proj_bundle, 4x the pair width) is materialized in full over (L, L, ...)
-    # regardless, which is what actually OOMs on large complexes. The fused cuequivariance
-    # kernel avoids materializing that intermediate at all; requires a CUDA 13 runtime,
-    # which we can't confirm ZeroGPU has without testing live, hence the fallback.
-    model.set_kernel_backend("cuequivariance")
-    print("Enabled cuequivariance fused kernels for triangle multiplication.")
-except Exception as exc:
-    print(f"cuequivariance backend unavailable ({exc}); falling back to the chunked path.")
-
-# `pair_loop_proj` (LayerNorm -> Linear, run once per recycling loop on the full (L, L, d_pair)
-# pair tensor) has no chunking in the checkpoint's own code, unlike pair_transition, which chunks
-# correctly. Both LayerNorm-over-last-dim and Linear are purely position-wise (no mixing across
-# different (i, j) pairs), so chunking along the sequence axis here is mathematically exact, not
-# an approximation.
-_PAIR_LOOP_PROJ_CHUNK_SIZE = 32
-_original_pair_loop_proj_forward = model.pair_loop_proj.forward
-
-
-def _chunked_pair_loop_proj_forward(x: torch.Tensor) -> torch.Tensor:
-    if x.shape[1] <= _PAIR_LOOP_PROJ_CHUNK_SIZE:
-        return _original_pair_loop_proj_forward(x)
-    out = torch.empty_like(x)
-    for start in range(0, x.shape[1], _PAIR_LOOP_PROJ_CHUNK_SIZE):
-        end = min(start + _PAIR_LOOP_PROJ_CHUNK_SIZE, x.shape[1])
-        out[:, start:end] = _original_pair_loop_proj_forward(x[:, start:end])
-    return out
-
-
-model.pair_loop_proj.forward = _chunked_pair_loop_proj_forward
-
-input_types = model.input_types
+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")
 print("Model loaded.")
 
 
@@ -76,83 +38,148 @@ def clean_sequence(raw_sequence: str) -> str:
     return re.sub(r"[^A-Z]", "", sequence)
 
 
-def render_structure(cif_str: str) -> str:
+def render_structure(pdb_str: str) -> str:
     view = py3Dmol.view(width=600, height=500)
-    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.addModel(pdb_str, "pdb")
+    view.setStyle({"cartoon": {"colorscheme": {"prop": "b", "gradient": "roygb", "min": 50, "max": 90}}})
     view.zoomTo()
     raw_html = view._make_html()
     escaped = raw_html.replace('"', "&quot;")
     return f'<iframe style="width: 100%; height: 500px; border: none;" srcdoc="{escaped}"></iframe>'
 
 
-@spaces.GPU(duration=300)
+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")
 def fold_rows(*args, progress=gr.Progress()):
     *seqs, num_loops, num_sampling_steps = args
 
-    chains: list[str] = []
-    for row_idx, (raw_seq, copies) in enumerate(zip(seqs[:NUM_ROWS], seqs[NUM_ROWS:])):
+    chain_sequences: list[str] = []
+    for raw_seq, copies in 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))}")
-        chains.extend([sequence] * max(1, int(copies or 1)))
+        chain_sequences.extend([sequence] * max(1, int(copies or 1)))
 
-    if not chains:
+    if not chain_sequences:
         raise gr.Error("Please enter at least one protein sequence.")
 
-    total_residues = sum(len(s) for s in chains)
+    total_residues = sum(len(s) for s in chain_sequences)
 
-    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)
-        ]
-    )
+    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")
 
     start = time.time()
-    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,
-        )
+    output = model.fold(
+        **features,
+        num_loops=num_loops,
+        num_sampling_steps=num_sampling_steps,
+        num_diffusion_samples=1,
+    )
     elapsed = time.time() - start
 
     progress(0.9, desc="Building structure file...")
-    cif_str = model.result_to_cif(result)
+    pdb_str = output_to_pdb(output, features)
 
-    cif_path = os.path.join(tempfile.gettempdir(), f"foldy_{int(time.time())}.cif")
-    with open(cif_path, "w") as f:
-        f.write(cif_str)
+    pdb_path = os.path.join(tempfile.gettempdir(), f"foldy_{int(time.time())}.pdb")
+    with open(pdb_path, "w") as f:
+        f.write(pdb_str)
 
-    viewer_html = render_structure(cif_str)
+    viewer_html = render_structure(pdb_str)
+    mean_plddt = float(output["plddt"].mean())
+    ptm = float(output["ptm"].mean())
     summary_lines = [
-        f"**Chains:** {len(chains)}",
+        f"**Chains:** {len(chain_sequences)}",
         f"**Total length:** {total_residues} residues",
-        "**Confidence scores:** unavailable (this checkpoint's confidence head is disabled)",
-        f"**Inference time:** {elapsed:.1f}s",
+        f"**Mean pLDDT:** {mean_plddt:.1f} / 100",
+        f"**pTM:** {ptm:.3f}",
     ]
+    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), cif_path
+    return viewer_html, "  \n".join(summary_lines), pdb_path
 
 
 with gr.Blocks(title="Foldy — ESMFold2 Protein Structure Prediction") as demo:
     gr.Markdown(
         """
         # 🧬 Foldy
-        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.
+        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.
         """
     )
 
@@ -179,13 +206,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=3, step=1, label="Recycling loops")
-                num_sampling_steps = gr.Slider(10, 50, value=15, step=5, label="Diffusion sampling steps")
+                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")
                 gr.Markdown("Higher values improve accuracy but increase inference time.")
 
             fold_btn = gr.Button("Fold", variant="primary")
             summary_output = gr.Markdown()
-            cif_output = gr.File(label="Download mmCIF", interactive=False)
+            pdb_output = gr.File(label="Download PDB", interactive=False)
 
         with gr.Column(scale=2):
             structure_output = gr.HTML(label="Predicted structure")
@@ -196,7 +223,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, cif_output],
+        outputs=[structure_output, summary_output, pdb_output],
     )
 
 if __name__ == "__main__":

requirements.txt

@@ -1,21 +1,2 @@
-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
+transformers>=5.16
 py3Dmol
-cuequivariance-torch
-cuequivariance-ops-torch-cu13