Use canonical .to('cuda') pattern + progress logs so container log shows what loader is doing

llm/zerogpu_backend.py

@@ -42,28 +42,43 @@ _tokenizer = None
 _load_done = threading.Event()
 
 
+def _log(msg: str) -> None:
+    """stdout flush so progress shows up in the HF container log immediately."""
+    print(f"[zerogpu_backend] {msg}", flush=True)
+
+
 def _load_model() -> None:
-    """Heavy CPU work (download ~16 GB safetensors + load to cuda emulation).
-    Runs in a daemon thread so that import returns immediately and the HTTP
-    server can serve /api/setup/status without blocking 15+ minutes."""
+    """Heavy CPU work (download ~16 GB safetensors + place on cuda). Runs in a
+    daemon thread so import returns immediately and the HTTP server can serve
+    /api/setup/status without blocking. Mirrors the canonical pattern from
+    https://huggingface.co/docs/hub/en/spaces-zerogpu :
+
+        pipe = X.from_pretrained(...)
+        pipe.to('cuda')
+    """
     global _model, _tokenizer, _LOAD_ERROR
     try:
+        _log(f"loading tokenizer for {_MODEL_ID}")
         _tokenizer = AutoTokenizer.from_pretrained(_MODEL_ID, trust_remote_code=True)
-        # On ZeroGPU, `device_map="cuda"` lands on the emulated device at import
-        # and becomes the real GPU inside @spaces.GPU. Off-platform (no `spaces`
-        # package installed), fall back to CPU so the module is still importable.
-        device_map = "cuda" if _ON_ZEROGPU else "cpu"
+        _log("tokenizer ready; loading model weights (this downloads ~16 GB on first run)")
         model = AutoModelForCausalLM.from_pretrained(
             _MODEL_ID,
             torch_dtype=torch.bfloat16,
             trust_remote_code=True,
-            device_map=device_map,
             low_cpu_mem_usage=True,
         )
+        _log("weights loaded to CPU; placing on cuda")
+        if _ON_ZEROGPU:
+            model = model.to("cuda")
+            _log("model on cuda (emulated until @spaces.GPU acquires real GPU)")
+        else:
+            _log("not on ZeroGPU; keeping model on CPU for local dev")
         model.eval()
         _model = model
+        _log("READY")
     except Exception as exc:  # pragma: no cover
         _LOAD_ERROR = f"{exc.__class__.__name__}: {exc}"
+        _log(f"LOAD FAILED: {_LOAD_ERROR}")
     finally:
         _load_done.set()