Preload ProtT5 weights on CPU before the GPU call, increase GPU duration

The first @spaces.GPU-decorated call was doing the full from_pretrained()
weight deserialization + device transfer inside the ZeroGPU time budget,
which could exceed the default duration and abort. Now the CPU-side load
happens eagerly at app startup, so the GPU call only does a fast device
transfer. Also set duration=180 as a safety margin.

app.py

@@ -33,6 +33,7 @@ AVAILABLE_EXAMPLES = [example for example in EXAMPLES if feature_path(example["i
 from t5_utils import (
     T5_REPO_ID,
     extract_t5_feature as _extract_t5_feature_core,
+    preload_t5_cpu,
     read_feature_h5,
 )
 
@@ -49,6 +50,13 @@ BASE_DIR = Path(__file__).parent
 # holding a ZeroGPU allocation (which has a short time budget).
 snapshot_download(T5_REPO_ID)
 
+# Also deserialize the weights into CPU memory here, still outside any
+# @spaces.GPU context. Without this, T5EncoderModel.from_pretrained() (loading
+# ~2.9GB from disk) would run lazily inside the first GPU-decorated call and
+# could eat enough of the ZeroGPU time budget to abort the task. With this,
+# the first GPU call only needs a fast .to("cuda") transfer.
+preload_t5_cpu()
+
 # ── Physical constants (from preprocess/misc.py) ──────────────────────────────
 MAX_TEMP   = 60.0
 MAX_CONC   = 1000.0
@@ -106,7 +114,7 @@ def load_llps_model():
 # features in assets/ always match what this would compute live. Only the
 # @spaces.GPU wrapping (ZeroGPU allocation) is app-specific, and cb_extract()
 # below skips calling this entirely when a cached feature is available.
-@spaces.GPU
+@spaces.GPU(duration=180)
 def extract_t5_feature(sequence: str) -> np.ndarray:
     return _extract_t5_feature_core(sequence)
 

t5_utils.py

@@ -24,45 +24,71 @@ def get_device() -> str:
     return "cuda" if torch.cuda.is_available() else "cpu"
 
 
-def load_t5():
+def preload_t5_cpu():
+    """Deserialize tokenizer + model weights onto CPU, outside any @spaces.GPU
+    context.
+
+    This is the slow step (reading ~2.9GB of weights off disk into RAM) and
+    it does not need a GPU. Calling it eagerly at app startup means the first
+    @spaces.GPU-decorated call only has to do a device transfer (fast), not a
+    full from_pretrained() load, so it no longer risks running past the
+    ZeroGPU time budget.
+    """
     global _t5_tokenizer, _t5_model
-    if _t5_model is None:
-        from transformers import AutoTokenizer, T5EncoderModel
-        import transformers.utils.import_utils as _hf_utils
-        import transformers.modeling_utils as _modeling_utils
-
-        # Rostlab/prot_t5_xl_half_uniref50-enc is only available as .bin (no
-        # safetensors). transformers 5.x blocks torch.load on torch < 2.6 due
-        # to CVE-2025-32434. We bypass that gate for this specific trusted
-        # checkpoint from the official HuggingFace Hub. The check lives in
-        # two places — import_utils AND the locally-imported name in
-        # modeling_utils — so both must be patched.
-        _noop = lambda: None
-        _orig_hf = _hf_utils.check_torch_load_is_safe
-        _orig_mdl = _modeling_utils.check_torch_load_is_safe
-
-        _hf_utils.check_torch_load_is_safe = _noop
-        _modeling_utils.check_torch_load_is_safe = _noop
-
-        try:
-            repo_id = T5_REPO_ID
-            dev = get_device()
-            dtype = torch.float16 if dev == "cuda" else torch.float32
-
-            _t5_tokenizer = AutoTokenizer.from_pretrained(
-                repo_id, do_lower_case=False, local_files_only=True
-            )
-            _t5_model = (
-                T5EncoderModel.from_pretrained(
-                    repo_id, torch_dtype=dtype, local_files_only=True
-                )
-                .to(dev)
-                .eval()
+    if _t5_model is not None:
+        return
+
+    from transformers import AutoTokenizer, T5EncoderModel
+    import transformers.utils.import_utils as _hf_utils
+    import transformers.modeling_utils as _modeling_utils
+
+    # Rostlab/prot_t5_xl_half_uniref50-enc is only available as .bin (no
+    # safetensors). transformers 5.x blocks torch.load on torch < 2.6 due
+    # to CVE-2025-32434. We bypass that gate for this specific trusted
+    # checkpoint from the official HuggingFace Hub. The check lives in
+    # two places — import_utils AND the locally-imported name in
+    # modeling_utils — so both must be patched.
+    _noop = lambda: None
+    _orig_hf = _hf_utils.check_torch_load_is_safe
+    _orig_mdl = _modeling_utils.check_torch_load_is_safe
+
+    _hf_utils.check_torch_load_is_safe = _noop
+    _modeling_utils.check_torch_load_is_safe = _noop
+
+    try:
+        repo_id = T5_REPO_ID
+        _t5_tokenizer = AutoTokenizer.from_pretrained(
+            repo_id, do_lower_case=False, local_files_only=True
+        )
+        _t5_model = (
+            T5EncoderModel.from_pretrained(
+                repo_id, torch_dtype=torch.float32, local_files_only=True
             )
-            _t5_model.requires_grad_(False)
-        finally:
-            _hf_utils.check_torch_load_is_safe = _orig_hf
-            _modeling_utils.check_torch_load_is_safe = _orig_mdl
+            .eval()
+        )
+        _t5_model.requires_grad_(False)
+    finally:
+        _hf_utils.check_torch_load_is_safe = _orig_hf
+        _modeling_utils.check_torch_load_is_safe = _orig_mdl
+
+
+def load_t5():
+    """Return (tokenizer, model) with the model on the current device.
+
+    Assumes preload_t5_cpu() has already been called (at app startup). The
+    only work left here is moving the already-deserialized model to the GPU
+    on the first @spaces.GPU call — a fast operation compared to loading it
+    from disk.
+    """
+    global _t5_model
+    preload_t5_cpu()
+
+    dev = get_device()
+    if next(_t5_model.parameters()).device.type != dev:
+        if dev == "cuda":
+            _t5_model = _t5_model.half().to(dev).eval()
+        else:
+            _t5_model = _t5_model.float().to(dev).eval()
 
     return _t5_tokenizer, _t5_model