Make model loading lazy at first inference to avoid startup crashes/timeouts

app.py

@@ -1,5 +1,6 @@
-import json
-from pathlib import Path
+import json
+import threading
+from pathlib import Path
 
 import gradio as gr
 import torch
@@ -9,8 +10,15 @@ from safetensors.torch import load_file
 from transformers import AutoImageProcessor, AutoModel
 from huggingface_hub import snapshot_download
 
-MODEL_REPO = "nzs234/siglip2-so400m-aesthetic-scorer-v1"
-CACHE_DIR = Path("./model_cache")
+MODEL_REPO = "nzs234/siglip2-so400m-aesthetic-scorer-v1"
+CACHE_DIR = Path("./model_cache")
+_STATE_LOCK = threading.Lock()
+_MODEL_READY = False
+_MODEL_ERR = ""
+processor = None
+model = None
+score_min = 1.0
+score_max = 9.0
 
 
 def infer_feature_dim(vision):
@@ -75,33 +83,51 @@ class Regressor(nn.Module):
         return torch.sigmoid(x)
 
 
-print("Downloading model repo snapshot...")
-local_repo = snapshot_download(repo_id=MODEL_REPO, repo_type="model", local_dir=str(CACHE_DIR))
-local_repo = Path(local_repo)
-meta = json.loads((local_repo / "metadata.json").read_text(encoding="utf-8"))
-model_cfg = meta.get("model", {})
-data_cfg = meta.get("data", {})
-
-processor = AutoImageProcessor.from_pretrained(str(local_repo / "backbone"), local_files_only=True, use_fast=False)
-model = Regressor(
-    backbone_dir=str(local_repo / "backbone"),
-    hidden_dim=int(model_cfg.get("hidden_dim", 2048)),
-    dropout=float(model_cfg.get("dropout", 0.2)),
-)
-head_state = load_file(str(local_repo / "head.safetensors"), device="cpu")
-model.head.load_state_dict(head_state, strict=False)
-model.eval()
-
-score_min = float(data_cfg.get("score_min", 1.0))
-score_max = float(data_cfg.get("score_max", 9.0))
+def _ensure_loaded():
+    global _MODEL_READY, _MODEL_ERR, processor, model, score_min, score_max
+    if _MODEL_READY:
+        return
+    with _STATE_LOCK:
+        if _MODEL_READY:
+            return
+        try:
+            print("Downloading model repo snapshot...")
+            local_repo = snapshot_download(repo_id=MODEL_REPO, repo_type="model", local_dir=str(CACHE_DIR))
+            local_repo = Path(local_repo)
+            meta = json.loads((local_repo / "metadata.json").read_text(encoding="utf-8"))
+            model_cfg = meta.get("model", {})
+            data_cfg = meta.get("data", {})
+
+            processor = AutoImageProcessor.from_pretrained(str(local_repo / "backbone"), local_files_only=True, use_fast=False)
+            model = Regressor(
+                backbone_dir=str(local_repo / "backbone"),
+                hidden_dim=int(model_cfg.get("hidden_dim", 2048)),
+                dropout=float(model_cfg.get("dropout", 0.2)),
+            )
+            head_state = load_file(str(local_repo / "head.safetensors"), device="cpu")
+            model.head.load_state_dict(head_state, strict=False)
+            model.eval()
+
+            score_min = float(data_cfg.get("score_min", 1.0))
+            score_max = float(data_cfg.get("score_max", 9.0))
+            _MODEL_READY = True
+            _MODEL_ERR = ""
+            print("Model loaded.")
+        except Exception as e:
+            _MODEL_ERR = str(e)
+            raise
 
 
 def predict(img: Image.Image):
     if img is None:
         return "error: no image"
-    if img.mode != "RGB":
-        img = img.convert("RGB")
-    proc = processor(images=img, return_tensors="pt")
+    try:
+        _ensure_loaded()
+    except Exception:
+        return f"error: model load failed: {_MODEL_ERR}"
+    if img.mode != "RGB":
+        img = img.convert("RGB")
+    proc = processor(images=img, return_tensors="pt")
     with torch.inference_mode():
         pred_01 = model(proc["pixel_values"]).item()
     pred_01 = max(0.0, min(1.0, float(pred_01)))