luh1124 Claude Sonnet 4.6 commited on
Commit
1aad7b7
Β·
1 Parent(s): c7a74a4

fix: replace DINOv2 torch.hub.load with download_url_to_file in preload

Browse files

torch.hub.load instantiates DINOv2 which imports xformers β†’ triggers CUDA
context init in the main process before any @GPU callback, breaking ZeroGPU.

Replace with torch.hub.download_url_to_file for the 1.13GB weights only.
The GPU callback still downloads the small GitHub repo code on cold start
but loads weights from local cache β€” no CUDA init in main process.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Files changed (1) hide show
  1. app.py +16 -10
app.py CHANGED
@@ -170,20 +170,26 @@ def _preload_worker() -> None:
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 ────────────────────────────────────────────────────────
 
170
  except Exception as exc:
171
  print(f"[NeAR] preload: RMBG-2.0 disk cache failed: {exc}", flush=True)
172
 
173
+ # Step 4: pre-download DINOv2 weights file only (no model instantiation).
174
+ # torch.hub.load instantiates the model which imports xformers β†’ triggers CUDA init
175
+ # in the main process, breaking ZeroGPU's context management.
176
+ # download_url_to_file is pure urllib β€” no CUDA. The GPU callback will still need
177
+ # to download the small GitHub repo code on cold start, but the 1.13 GB weights
178
+ # file is the slow part and will be served from this local cache.
179
  if not (os.environ.get("NEAR_DINO_LOCAL_REPO") or os.environ.get("NEAR_AUX_REPO")):
180
  try:
181
  import torch
182
+ ckpt_dir = os.path.join(torch.hub.get_dir(), "checkpoints")
183
+ os.makedirs(ckpt_dir, exist_ok=True)
184
+ ckpt_path = os.path.join(ckpt_dir, "dinov2_vitl14_reg4_pretrain.pth")
185
+ if not os.path.exists(ckpt_path):
186
+ torch.hub.download_url_to_file(
187
+ "https://dl.fbaipublicfiles.com/dinov2/dinov2_vitl14/dinov2_vitl14_reg4_pretrain.pth",
188
+ ckpt_path, progress=True,
189
+ )
190
+ print("[NeAR] preload: DINOv2 weights file cached.", flush=True)
191
  except Exception as exc:
192
+ print(f"[NeAR] preload: DINOv2 weight prefetch failed: {exc}", flush=True)
193
 
194
 
195
  # ── GPU ensure helpers ────────────────────────────────────────────────────────