Add ESMFold folding API

.dockerignore

@@ -0,0 +1,6 @@
+__pycache__/
+*.pyc
+.pytest_cache/
+.ruff_cache/
+tests/
+

Dockerfile

@@ -0,0 +1,26 @@
+FROM pytorch/pytorch:2.5.1-cuda12.4-cudnn9-runtime
+
+ENV PYTHONDONTWRITEBYTECODE=1 \
+    PYTHONUNBUFFERED=1 \
+    HF_HOME=/data/huggingface \
+    TRANSFORMERS_CACHE=/data/huggingface \
+    FOLD_BACKEND=esmfold \
+    MAX_PROTEIN_AA=400
+
+WORKDIR /app
+
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends git git-lfs && \
+    rm -rf /var/lib/apt/lists/* && \
+    git lfs install
+
+COPY requirements.txt /app/requirements.txt
+RUN pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir -r /app/requirements.txt
+
+COPY . /app
+
+EXPOSE 7860
+
+CMD ["uvicorn", "folding_api_service.app:app", "--host", "0.0.0.0", "--port", "7860"]
+

README.md

@@ -1,10 +1,45 @@
 ---
-title: Protein Folding Api
-emoji: 🌖
-colorFrom: blue
-colorTo: pink
+title: Carbon Protein Folding API
+emoji: 🧬
+colorFrom: green
+colorTo: blue
 sdk: docker
-pinned: false
+app_port: 7860
+fullWidth: true
+tags:
+  - biology
+  - protein-folding
+  - esmfold
+  - fastapi
+  - carbon
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Carbon Protein Folding API
+
+FastAPI service for the Carbon DNA-to-Structure demo. The first live backend is ESMFold for single-chain protein folding.
+
+## Endpoints
+
+- `GET /health`
+- `GET /tools`
+- `POST /jobs`
+- `GET /jobs/{job_id}`
+
+`POST /jobs` accepts one protein entity and returns immediately with a `job_id`. Poll `GET /jobs/{job_id}` until the job reaches `succeeded` or `failed`.
+
+## Configuration
+
+Set these Space variables/secrets:
+
+```sh
+FOLD_BACKEND=esmfold
+FOLD_API_TOKEN=...
+MAX_PROTEIN_AA=400
+```
+
+For local CPU tests, use:
+
+```sh
+FOLD_BACKEND=stub uvicorn folding_api_service.app:app --host 0.0.0.0 --port 7860
+```
+

app.py

@@ -0,0 +1,2 @@
+from folding_api_service.app import app
+

folding_api_service/__init__.py

@@ -0,0 +1,2 @@
+"""Protein folding API service for the Carbon demo."""
+

folding_api_service/app.py

@@ -0,0 +1,198 @@
+from __future__ import annotations
+
+import os
+import re
+import time
+import uuid
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import asdict, dataclass, field
+from threading import Lock
+from typing import Any
+
+from fastapi import Depends, FastAPI, Header, HTTPException
+from pydantic import BaseModel, Field
+
+from .backends import FoldingBackend, FoldOutput, make_backend
+
+
+PROTEIN_RE = re.compile(r"^[ACDEFGHIKLMNPQRSTVWYXBZUOJ]+$", re.IGNORECASE)
+MAX_PROTEIN_AA = int(os.getenv("MAX_PROTEIN_AA", "400"))
+API_TOKEN = os.getenv("FOLD_API_TOKEN", "").strip()
+
+
+class Entity(BaseModel):
+    id: str = Field(min_length=1, max_length=32)
+    type: str
+    sequence: str
+
+
+class JobRequest(BaseModel):
+    tool_id: str
+    entities: list[Entity]
+    options: dict[str, Any] = Field(default_factory=dict)
+    client_metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+@dataclass
+class JobState:
+    job_id: str
+    tool_id: str
+    status: str
+    created_at: float
+    updated_at: float
+    progress: float = 0.0
+    result: dict[str, Any] | None = None
+    error: str | None = None
+
+    def public(self) -> dict[str, Any]:
+        payload = asdict(self)
+        payload.pop("created_at", None)
+        payload.pop("updated_at", None)
+        return payload
+
+
+@dataclass
+class RuntimeState:
+    backend: FoldingBackend = field(default_factory=make_backend)
+    jobs: dict[str, JobState] = field(default_factory=dict)
+    lock: Lock = field(default_factory=Lock)
+    executor: ThreadPoolExecutor = field(default_factory=lambda: ThreadPoolExecutor(max_workers=1))
+
+
+state = RuntimeState()
+app = FastAPI(title="Carbon Protein Folding API", version="0.1.0")
+
+
+def require_auth(authorization: str | None = Header(default=None)) -> None:
+    if not API_TOKEN:
+        return
+    expected = f"Bearer {API_TOKEN}"
+    if authorization != expected:
+        raise HTTPException(status_code=401, detail="invalid or missing bearer token")
+
+
+def validate_request(payload: JobRequest) -> str:
+    if payload.tool_id != "esmfold":
+        raise HTTPException(status_code=400, detail="only tool_id 'esmfold' is supported")
+    if len(payload.entities) != 1:
+        raise HTTPException(status_code=400, detail="exactly one protein entity is supported")
+
+    entity = payload.entities[0]
+    if entity.type.lower() != "protein":
+        raise HTTPException(status_code=400, detail="entity type must be 'protein'")
+
+    sequence = re.sub(r"\s+", "", entity.sequence).upper().replace("*", "")
+    if not sequence:
+        raise HTTPException(status_code=400, detail="protein sequence is empty")
+    if len(sequence) > MAX_PROTEIN_AA:
+        raise HTTPException(status_code=400, detail=f"protein sequence exceeds {MAX_PROTEIN_AA} aa")
+    if not PROTEIN_RE.match(sequence):
+        raise HTTPException(status_code=400, detail="protein sequence contains unsupported characters")
+    return sequence
+
+
+@app.get("/health")
+def health() -> dict[str, Any]:
+    return {
+        "ok": True,
+        "backend": os.getenv("FOLD_BACKEND", "esmfold"),
+        "max_protein_aa": MAX_PROTEIN_AA,
+        "jobs": len(state.jobs),
+    }
+
+
+@app.get("/tools")
+def tools(_: None = Depends(require_auth)) -> dict[str, Any]:
+    return {
+        "tools": [
+            {
+                "id": "esmfold",
+                "name": "ESMFold",
+                "status": "live",
+                "input_types": ["protein"],
+                "max_protein_aa": MAX_PROTEIN_AA,
+                "output_formats": ["pdb"],
+                "options": {
+                    "seed": {"type": "integer", "supported": False},
+                    "num_recycles": {"type": "integer", "supported": False},
+                    "msa_mode": {"type": "string", "value": "none"},
+                },
+            }
+        ]
+    }
+
+
+@app.post("/jobs")
+def create_job(payload: JobRequest, _: None = Depends(require_auth)) -> dict[str, str]:
+    sequence = validate_request(payload)
+    job_id = uuid.uuid4().hex
+    now = time.time()
+    job = JobState(
+        job_id=job_id,
+        tool_id=payload.tool_id,
+        status="queued",
+        created_at=now,
+        updated_at=now,
+    )
+    with state.lock:
+        state.jobs[job_id] = job
+    state.executor.submit(run_job, job_id, sequence, payload.options)
+    return {"job_id": job_id, "status": "queued"}
+
+
+@app.get("/jobs/{job_id}")
+def get_job(job_id: str, _: None = Depends(require_auth)) -> dict[str, Any]:
+    with state.lock:
+        job = state.jobs.get(job_id)
+        if job is None:
+            raise HTTPException(status_code=404, detail="unknown job_id")
+        return job.public()
+
+
+def run_job(job_id: str, sequence: str, options: dict[str, Any]) -> None:
+    _update_job(job_id, status="running", progress=0.05)
+    try:
+        output = state.backend.fold(sequence, options)
+        _update_job(
+            job_id,
+            status="succeeded",
+            progress=1.0,
+            result=_result_payload(output),
+            error=None,
+        )
+    except Exception as exc:  # noqa: BLE001 - API should preserve job failure details.
+        _update_job(job_id, status="failed", progress=1.0, error=str(exc), result=None)
+
+
+def _update_job(
+    job_id: str,
+    *,
+    status: str,
+    progress: float,
+    result: dict[str, Any] | None = None,
+    error: str | None = None,
+) -> None:
+    with state.lock:
+        job = state.jobs[job_id]
+        job.status = status
+        job.progress = progress
+        job.updated_at = time.time()
+        if result is not None:
+            job.result = result
+        if error is not None:
+            job.error = error
+
+
+def _result_payload(output: FoldOutput) -> dict[str, Any]:
+    return {
+        "structures": [
+            {
+                "format": "pdb",
+                "content": output.pdb,
+                "confidence": output.confidence,
+            }
+        ],
+        "metrics": output.metrics,
+        "warnings": output.warnings,
+    }
+

folding_api_service/backends.py

@@ -0,0 +1,181 @@
+from __future__ import annotations
+
+import os
+import time
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any
+
+
+DEMO_PDB = """HEADER    CARBON FOLDING API STUB
+ATOM      1  N   ALA A   1      -0.500   1.300   0.000  1.00 80.00           N
+ATOM      2  CA  ALA A   1       0.000   0.000   0.000  1.00 80.00           C
+ATOM      3  C   ALA A   1       1.520   0.000   0.000  1.00 80.00           C
+ATOM      4  O   ALA A   1       2.110  -1.060   0.000  1.00 80.00           O
+ATOM      5  N   GLY A   2       2.160   1.170   0.000  1.00 82.00           N
+ATOM      6  CA  GLY A   2       3.600   1.260   0.000  1.00 82.00           C
+ATOM      7  C   GLY A   2       4.160   2.660   0.000  1.00 82.00           C
+ATOM      8  O   GLY A   2       3.480   3.660   0.000  1.00 82.00           O
+ATOM      9  N   SER A   3       5.430   2.730   0.000  1.00 76.00           N
+ATOM     10  CA  SER A   3       6.080   4.030   0.000  1.00 76.00           C
+ATOM     11  C   SER A   3       7.600   3.910   0.000  1.00 76.00           C
+ATOM     12  O   SER A   3       8.250   4.920   0.000  1.00 76.00           O
+TER
+END
+"""
+
+
+@dataclass(frozen=True)
+class FoldOutput:
+    pdb: str
+    confidence: dict[str, Any]
+    metrics: dict[str, Any]
+    warnings: list[str]
+
+
+class FoldingBackend(ABC):
+    @abstractmethod
+    def fold(self, sequence: str, options: dict[str, Any]) -> FoldOutput:
+        raise NotImplementedError
+
+
+class StubBackend(FoldingBackend):
+    def fold(self, sequence: str, options: dict[str, Any]) -> FoldOutput:
+        del options
+        started = time.monotonic()
+        time.sleep(min(0.1, max(0.0, len(sequence) / 10_000)))
+        return FoldOutput(
+            pdb=DEMO_PDB,
+            confidence={"mean_plddt": 80.0},
+            metrics={"runtime_seconds": round(time.monotonic() - started, 4), "sequence_length": len(sequence)},
+            warnings=["stub backend returned a demo structure"],
+        )
+
+
+class EsmFoldBackend(FoldingBackend):
+    def __init__(self, model_id: str = "facebook/esmfold_v1") -> None:
+        self.model_id = model_id
+        self._loaded = False
+        self._device = None
+        self._tokenizer = None
+        self._model = None
+
+    def _load(self) -> None:
+        if self._loaded:
+            return
+
+        import torch
+        from transformers import AutoTokenizer, EsmForProteinFolding
+
+        self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self._tokenizer = AutoTokenizer.from_pretrained(self.model_id)
+        self._model = EsmForProteinFolding.from_pretrained(
+            self.model_id,
+            low_cpu_mem_usage=True,
+        )
+        self._model.eval()
+        self._model.to(self._device)
+
+        # Reduce memory use for longer demo proteins. This is supported by the
+        # Transformers ESMFold implementation and is a no-op if unavailable.
+        if hasattr(self._model, "trunk") and hasattr(self._model.trunk, "set_chunk_size"):
+            self._model.trunk.set_chunk_size(int(os.getenv("ESMFOLD_CHUNK_SIZE", "64")))
+
+        self._loaded = True
+
+    def fold(self, sequence: str, options: dict[str, Any]) -> FoldOutput:
+        del options
+        started = time.monotonic()
+        self._load()
+
+        import torch
+
+        assert self._device is not None
+        assert self._tokenizer is not None
+        assert self._model is not None
+
+        tokenized = self._tokenizer([sequence], return_tensors="pt", add_special_tokens=False)
+        tokenized = {key: value.to(self._device) for key, value in tokenized.items()}
+
+        with torch.no_grad():
+            output = self._model(**tokenized)
+
+        pdb = _esmfold_output_to_pdb(output)
+        mean_plddt = _mean_plddt(output)
+        runtime = time.monotonic() - started
+        warnings = []
+        if self._device.type != "cuda":
+            warnings.append("ESMFold ran on CPU; GPU is recommended")
+        if mean_plddt is not None and mean_plddt < 50:
+            warnings.append("low mean pLDDT; predicted structure may be unreliable")
+
+        return FoldOutput(
+            pdb=pdb,
+            confidence={"mean_plddt": mean_plddt},
+            metrics={
+                "runtime_seconds": round(runtime, 4),
+                "sequence_length": len(sequence),
+                "device": self._device.type,
+            },
+            warnings=warnings,
+        )
+
+
+def _as_mapping(output: Any) -> dict[str, Any]:
+    if isinstance(output, dict):
+        return output
+    if hasattr(output, "to_tuple") and hasattr(output, "keys"):
+        return {key: output[key] for key in output.keys()}
+    if hasattr(output, "__dict__"):
+        return {key: value for key, value in vars(output).items() if not key.startswith("_")}
+    raise TypeError("unsupported ESMFold output type")
+
+
+def _esmfold_output_to_pdb(output: Any) -> str:
+    import torch
+    from transformers.models.esm.openfold_utils.feats import atom14_to_atom37
+    from transformers.models.esm.openfold_utils.protein import Protein as OpenFoldProtein
+    from transformers.models.esm.openfold_utils.protein import to_pdb
+
+    data = _as_mapping(output)
+    final_atom_positions = atom14_to_atom37(data["positions"][-1], data)
+
+    cpu_data = {}
+    for key, value in data.items():
+        if torch.is_tensor(value):
+            cpu_data[key] = value.detach().cpu().numpy()
+        else:
+            cpu_data[key] = value
+
+    final_atom_positions = final_atom_positions.detach().cpu().numpy()
+    final_atom_mask = cpu_data["atom37_atom_exists"]
+
+    protein = OpenFoldProtein(
+        aatype=cpu_data["aatype"][0],
+        atom_positions=final_atom_positions[0],
+        atom_mask=final_atom_mask[0],
+        residue_index=cpu_data["residue_index"][0] + 1,
+        b_factors=cpu_data["plddt"][0],
+        chain_index=cpu_data.get("chain_index", [None])[0],
+    )
+    return to_pdb(protein)
+
+
+def _mean_plddt(output: Any) -> float | None:
+    data = _as_mapping(output)
+    plddt = data.get("plddt")
+    if plddt is None:
+        return None
+    if hasattr(plddt, "detach"):
+        return round(float(plddt.detach().float().mean().cpu().item()), 4)
+    return round(float(plddt.mean()), 4)
+
+
+def make_backend() -> FoldingBackend:
+    backend = os.getenv("FOLD_BACKEND", "esmfold").strip().lower()
+    if backend == "stub":
+        return StubBackend()
+    if backend == "esmfold":
+        return EsmFoldBackend(os.getenv("ESMFOLD_MODEL_ID", "facebook/esmfold_v1"))
+    raise ValueError(f"unsupported FOLD_BACKEND: {backend}")
+

requirements.txt

@@ -0,0 +1,9 @@
+accelerate>=0.31
+fastapi>=0.110
+huggingface-hub>=0.23
+numpy<2
+pydantic>=2.7
+scipy>=1.11
+transformers>=4.44,<5
+uvicorn[standard]>=0.29
+

tests/test_api.py

@@ -0,0 +1,76 @@
+import os
+import time
+import unittest
+
+os.environ.setdefault("FOLD_BACKEND", "stub")
+os.environ.setdefault("MAX_PROTEIN_AA", "400")
+
+from fastapi.testclient import TestClient  # noqa: E402
+
+from folding_api_service.app import app  # noqa: E402
+
+
+class FoldingApiTest(unittest.TestCase):
+    def setUp(self):
+        self.client = TestClient(app)
+
+    def test_health(self):
+        response = self.client.get("/health")
+        self.assertEqual(response.status_code, 200)
+        self.assertTrue(response.json()["ok"])
+
+    def test_tools(self):
+        response = self.client.get("/tools")
+        self.assertEqual(response.status_code, 200)
+        tools = response.json()["tools"]
+        self.assertEqual(tools[0]["id"], "esmfold")
+        self.assertEqual(tools[0]["max_protein_aa"], 400)
+
+    def test_submit_and_poll_stub_job(self):
+        response = self.client.post(
+            "/jobs",
+            json={
+                "tool_id": "esmfold",
+                "entities": [{"id": "A", "type": "protein", "sequence": "MLSDEDFKAVFGMTRSAFANLPLWKQQNLKKEKGLF"}],
+                "options": {"msa_mode": "none"},
+            },
+        )
+        self.assertEqual(response.status_code, 200)
+        job_id = response.json()["job_id"]
+
+        final = None
+        for _ in range(20):
+            poll = self.client.get(f"/jobs/{job_id}")
+            self.assertEqual(poll.status_code, 200)
+            final = poll.json()
+            if final["status"] in {"succeeded", "failed"}:
+                break
+            time.sleep(0.05)
+
+        self.assertIsNotNone(final)
+        self.assertEqual(final["status"], "succeeded")
+        self.assertIn("ATOM", final["result"]["structures"][0]["content"])
+        self.assertEqual(final["result"]["structures"][0]["format"], "pdb")
+
+    def test_validation_errors(self):
+        cases = [
+            {"tool_id": "omegafold", "entities": [{"id": "A", "type": "protein", "sequence": "MKT"}]},
+            {"tool_id": "esmfold", "entities": []},
+            {"tool_id": "esmfold", "entities": [{"id": "A", "type": "dna", "sequence": "ATG"}]},
+            {"tool_id": "esmfold", "entities": [{"id": "A", "type": "protein", "sequence": ""}]},
+            {"tool_id": "esmfold", "entities": [{"id": "A", "type": "protein", "sequence": "M1T"}]},
+            {"tool_id": "esmfold", "entities": [{"id": "A", "type": "protein", "sequence": "M" * 401}]},
+        ]
+        for payload in cases:
+            with self.subTest(payload=payload):
+                response = self.client.post("/jobs", json={**payload, "options": {}})
+                self.assertEqual(response.status_code, 400)
+
+    def test_unknown_job(self):
+        response = self.client.get("/jobs/missing")
+        self.assertEqual(response.status_code, 404)
+
+
+if __name__ == "__main__":
+    unittest.main()
+