luh1124 Claude Sonnet 4.6 commited on
Commit
c0d618a
Β·
1 Parent(s): 487f993

fix(preload): split into disk-cache phase + CPU-load phase to unblock GPU callbacks

Browse files

Previous: _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>

Files changed (1) hide show
  1. app.py +27 -6
app.py CHANGED
@@ -135,16 +135,37 @@ def _load_geometry_cpu_locked() -> None:
135
 
136
 
137
  def _preload_worker() -> None:
138
- # Only preload Hunyuan3D on CPU β€” it is safe (no CUDA init during loading).
139
- # NeAR triggers CUDA init during from_pretrained (BiRefNet/DINOv2 init with
140
- # trust_remote_code), which is not allowed in ZeroGPU's main process.
141
- # NeAR loads lazily inside the first @GPU callback instead.
 
 
 
 
 
 
142
  try:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
143
  with _MODEL_LOCK:
144
  _load_geometry_cpu_locked()
145
- print("[NeAR] CPU preload complete (Hunyuan3D only).", flush=True)
146
  except Exception as exc:
147
- print(f"[NeAR] CPU preload failed: {exc}", flush=True)
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 ────────────────────────────────────────────────────────