fix: replace DINOv2 torch.hub.load with download_url_to_file in preload
Browse filestorch.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>
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:
|
| 174 |
-
#
|
| 175 |
-
#
|
|
|
|
|
|
|
|
|
|
| 176 |
if not (os.environ.get("NEAR_DINO_LOCAL_REPO") or os.environ.get("NEAR_AUX_REPO")):
|
| 177 |
try:
|
| 178 |
import torch
|
| 179 |
-
|
| 180 |
-
|
| 181 |
-
|
| 182 |
-
)
|
| 183 |
-
|
| 184 |
-
|
|
|
|
|
|
|
|
|
|
| 185 |
except Exception as exc:
|
| 186 |
-
print(f"[NeAR] preload: DINOv2
|
| 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 ββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|