update

Dockerfile

@@ -17,6 +17,7 @@ ARG MAMBA_DOCKERFILE_ACTIVATE=1
 ENV BASH_ENV=/usr/local/bin/_activate_current_env.sh
 ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/conda/lib
 
+RUN pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cpu
 RUN pip install --no-cache-dir \
     "stoic @ git+https://github.com/PickyBinders/stoic.git" \
     gradio==6.9.0

app.py

@@ -1,3 +1,4 @@
+import functools
 import time
 
 import torch
@@ -6,11 +7,15 @@ from loguru import logger
 
 from stoic.model import Stoic
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-logger.info(f"Loading model on {device}")
-model = Stoic.from_pretrained("PickyBinders/stoic")
-model = model.to(device).eval()
-logger.info("Model loaded")
+
+@functools.lru_cache(maxsize=1)
+def get_model():
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    logger.info(f"Loading model on {device}")
+    model = Stoic.from_pretrained("PickyBinders/stoic")
+    model = model.to(device).eval()
+    logger.info("Model loaded")
+    return model
 
 
 def predict(sequences_text: str, top_n: int) -> tuple[str, str]:
@@ -20,6 +25,7 @@ def predict(sequences_text: str, top_n: int) -> tuple[str, str]:
     if len(sequences) > 26:
         raise gr.Error("Maximum 26 unique chains supported.")
 
+    model = get_model()
     start = time.time()
     with torch.no_grad():
         results = model.predict_stoichiometry(sequences, top_n=top_n)