fix: pre-cache RMBG-2.0 and DINOv2 in preload worker to prevent GPU lease timeout
Browse filesbriaai/RMBG-2.0 and DINOv2 (facebookresearch/dinov2) were being downloaded
inside the 240s @GPU callback on first use, exhausting the ZeroGPU lease before
inference could run. Preload worker now warms both caches at Space startup.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
app.py
CHANGED
|
@@ -161,6 +161,30 @@ def _preload_worker() -> None:
|
|
| 161 |
except Exception as exc:
|
| 162 |
print(f"[NeAR] preload: NeAR disk cache failed: {exc}", flush=True)
|
| 163 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 164 |
|
| 165 |
# ββ GPU ensure helpers ββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 166 |
# Called at the top of EVERY @GPU callback. Always re-creates renderer and
|
|
|
|
| 161 |
except Exception as exc:
|
| 162 |
print(f"[NeAR] preload: NeAR disk cache failed: {exc}", flush=True)
|
| 163 |
|
| 164 |
+
# Step 3: warm rembg model cache (briaai/RMBG-2.0, referenced in pipeline.yaml).
|
| 165 |
+
# Without this, the download happens inside the 240s GPU callback and times out.
|
| 166 |
+
try:
|
| 167 |
+
from huggingface_hub import snapshot_download
|
| 168 |
+
snapshot_download(repo_id="briaai/RMBG-2.0", token=os.environ.get("HF_TOKEN"))
|
| 169 |
+
print("[NeAR] preload: RMBG-2.0 disk cache ready.", flush=True)
|
| 170 |
+
except Exception as exc:
|
| 171 |
+
print(f"[NeAR] preload: RMBG-2.0 disk cache failed: {exc}", flush=True)
|
| 172 |
+
|
| 173 |
+
# Step 4: warm DINOv2 torch.hub cache.
|
| 174 |
+
# If NEAR_AUX_REPO is set, snapshot_download handles it inside load_dinov2_model.
|
| 175 |
+
# Otherwise we must pre-fetch facebookresearch/dinov2 from GitHub now (CPU-only).
|
| 176 |
+
if not (os.environ.get("NEAR_DINO_LOCAL_REPO") or os.environ.get("NEAR_AUX_REPO")):
|
| 177 |
+
try:
|
| 178 |
+
import torch
|
| 179 |
+
_dino_tmp = torch.hub.load(
|
| 180 |
+
"facebookresearch/dinov2", "dinov2_vitl14_reg",
|
| 181 |
+
pretrained=True, verbose=False,
|
| 182 |
+
)
|
| 183 |
+
del _dino_tmp
|
| 184 |
+
print("[NeAR] preload: DINOv2 torch.hub cache ready.", flush=True)
|
| 185 |
+
except Exception as exc:
|
| 186 |
+
print(f"[NeAR] preload: DINOv2 torch.hub cache failed: {exc}", flush=True)
|
| 187 |
+
|
| 188 |
|
| 189 |
# ββ GPU ensure helpers ββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 190 |
# Called at the top of EVERY @GPU callback. Always re-creates renderer and
|