luh1124 Claude Sonnet 4.6 commited on
Commit
76644b1
Β·
1 Parent(s): 914fb3d

fix: pre-cache RMBG-2.0 and DINOv2 in preload worker to prevent GPU lease timeout

Browse files

briaai/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>

Files changed (1) hide show
  1. app.py +24 -0
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