Pin torch 2.8.0 for ZeroGPU; build models on CPU at startup

app.py

@@ -229,5 +229,12 @@ with gr.Blocks(title="CELL-FM CondenSeq", theme=gr.themes.Soft()) as demo:
     )
 
 
+# Download weights and build the models at startup, on CPU. Under ZeroGPU the GPU
+# budget only covers the decorated call, so this must not happen inside it.
+try:
+    pipeline.load_models()
+except Exception as exc:  # surfaced in the Space logs; the UI still loads
+    print(f"[startup] model preload failed: {exc}")
+
 if __name__ == "__main__":
     demo.queue().launch(share=os.environ.get("GRADIO_SHARE") == "1")

pipeline.py

@@ -134,20 +134,29 @@ def get_device() -> torch.device:
 
 @lru_cache(maxsize=1)
 def load_models():
-    """Build both models once and keep them resident."""
-    device = get_device()
+    """Download the weights and build both models, on CPU.
 
+    Kept off the GPU deliberately: under ZeroGPU a CUDA device only exists inside a
+    @spaces.GPU call, and the 2.1 GB download plus model construction would eat that
+    budget. Call this once at startup, then models_on(device) inside the GPU call.
+    """
     generator = CELLFMCSModel(config=generator_config(), loss_fn=UniDiffCriterions)
-    generator.to(device).eval()
+    generator.eval()
 
     classifier = ViTModel(config=classifier_config())
-    classifier.to(device).eval()
+    classifier.eval()
 
     from esm.tokenization.sequence_tokenizer import EsmSequenceTokenizer
 
     return generator, classifier, EsmSequenceTokenizer()
 
 
+def models_on(device):
+    """Move the cached models to device. A no-op once they are already there."""
+    generator, classifier, vocab = load_models()
+    return generator.to(device), classifier.to(device), vocab
+
+
 @lru_cache(maxsize=1)
 def reference_nucleus() -> torch.Tensor:
     """The DAPI channel every generated image is conditioned on, as (1, 1, 160, 160) in [-1, 1].
@@ -208,8 +217,8 @@ def generate_images(sequence, num_images, num_steps, batch_size, seed=DEFAULT_SE
     and classification, so the classifier saw quantised pixels. Keeping the same
     round-trip here keeps demo numbers comparable to the published ones.
     """
-    generator, _, vocab = load_models()
     device = get_device()
+    generator, _, vocab = models_on(device)
 
     from esm.utils import encoding
 
@@ -248,8 +257,8 @@ def generate_images(sequence, num_images, num_steps, batch_size, seed=DEFAULT_SE
 @torch.no_grad()
 def classify_images(images, batch_size, progress=None):
     """Stage 2: condensed (1) vs diffuse (0) for each generated image."""
-    _, classifier, _ = load_models()
     device = get_device()
+    _, classifier, _ = models_on(device)
 
     preds = []
     for start in range(0, images.shape[0], batch_size):

requirements.txt

@@ -1,6 +1,9 @@
-# Pinned to the versions the CELL-FM checkpoints were trained and evaluated with.
-torch==2.4.1
-torchvision==0.19.1
+# ZeroGPU only accepts torch 2.8.0+; the checkpoints were trained on 2.4.1, and
+# state dicts are just tensors, so they load unchanged. torchvision is pinned to
+# the matching release for 2.8.0.
+torch==2.8.0
+torchvision==0.23.0
+spaces
 diffusers==0.31.0
 transformers==4.44.2
 esm==3.1.4