fix(preload): split into disk-cache phase + CPU-load phase to unblock GPU callbacks
Browse filesPrevious: _MODEL_LOCK held during entire snapshot_download (minutes) β GPU
callback blocks waiting for lock β ZeroGPU 240 s budget expires β error.
Fix: Phase 1 downloads both repos to HF disk cache without holding any lock.
Phase 2 instantiates Hunyuan3D from cache under lock (fast, no network I/O).
GPU callbacks can now acquire the lock quickly even if the user clicks early.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
app.py
CHANGED
|
@@ -135,16 +135,37 @@ def _load_geometry_cpu_locked() -> None:
|
|
| 135 |
|
| 136 |
|
| 137 |
def _preload_worker() -> None:
|
| 138 |
-
|
| 139 |
-
|
| 140 |
-
|
| 141 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 142 |
try:
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 143 |
with _MODEL_LOCK:
|
| 144 |
_load_geometry_cpu_locked()
|
| 145 |
-
print("[NeAR]
|
| 146 |
except Exception as exc:
|
| 147 |
-
print(f"[NeAR]
|
| 148 |
|
| 149 |
|
| 150 |
# ββ GPU ensure helpers ββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
|
|
|
| 135 |
|
| 136 |
|
| 137 |
def _preload_worker() -> None:
|
| 138 |
+
"""Two-phase background preload (avoids blocking GPU callbacks):
|
| 139 |
+
|
| 140 |
+
Phase 1 β disk cache warm-up (no lock, no CUDA):
|
| 141 |
+
snapshot_download both repos so GPU callbacks never wait for network I/O.
|
| 142 |
+
Phase 2 β CPU instantiation (lock held briefly, no CUDA):
|
| 143 |
+
Load Hunyuan3D weights into CPU RAM from the local cache.
|
| 144 |
+
|
| 145 |
+
NeAR instantiates lazily inside @GPU callbacks β its BiRefNet/DINOv2 init
|
| 146 |
+
triggers CUDA which is not allowed in the ZeroGPU main process.
|
| 147 |
+
"""
|
| 148 |
try:
|
| 149 |
+
from huggingface_hub import snapshot_download
|
| 150 |
+
|
| 151 |
+
# Phase 1: download repos to disk β no lock needed.
|
| 152 |
+
for _repo in (
|
| 153 |
+
os.environ.get("NEAR_HUNYUAN_PRETRAINED", "tencent/Hunyuan3D-2.1"),
|
| 154 |
+
"luh0502/NeAR",
|
| 155 |
+
):
|
| 156 |
+
try:
|
| 157 |
+
print(f"[NeAR] preload: caching {_repo!r} to diskβ¦", flush=True)
|
| 158 |
+
snapshot_download(repo_id=_repo, token=os.environ.get("HF_TOKEN"))
|
| 159 |
+
print(f"[NeAR] preload: {_repo!r} disk-cache ready.", flush=True)
|
| 160 |
+
except Exception as _exc:
|
| 161 |
+
print(f"[NeAR] preload: disk-cache {_repo!r} failed: {_exc}", flush=True)
|
| 162 |
+
|
| 163 |
+
# Phase 2: instantiate Hunyuan3D on CPU (fast now that files are on disk).
|
| 164 |
with _MODEL_LOCK:
|
| 165 |
_load_geometry_cpu_locked()
|
| 166 |
+
print("[NeAR] preload: complete.", flush=True)
|
| 167 |
except Exception as exc:
|
| 168 |
+
print(f"[NeAR] preload failed: {exc}", flush=True)
|
| 169 |
|
| 170 |
|
| 171 |
# ββ GPU ensure helpers ββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|