Load model on cuda at module level (canonical ZeroGPU pattern)

README.md

@@ -43,15 +43,16 @@ in a custom `gr.Server` HTML/JS board interface (no stock Gradio components).
 
 - **Model:** [`openbmb/MiniCPM4.1-8B`](https://huggingface.co/openbmb/MiniCPM4.1-8B)
   (text, bf16 transformers) — ~8B params, well under the 32B cap.
-- **Inference:** in-process Hugging Face `transformers`. The model lives in CPU
-  RAM and is moved to cuda inside a `@spaces.GPU`-decorated function, so the
-  H200 is only held while the model is actually generating.
+- **Inference:** in-process Hugging Face `transformers`, placed on `cuda` at
+  module load (using ZeroGPU's PyTorch CUDA emulation), with the real GPU
+  attached only inside a `@spaces.GPU`-decorated `generate()` call.
 
 ## 🖥️ Hardware
 
-Runs on **ZeroGPU** (H200, free 40 min/day for org members). Each generation is
-capped at `PHANTOM_GRID_ZEROGPU_DURATION` seconds (default 90). No voice path
-on this Space (`PHANTOM_GRID_WITNESS_CHAT_TTS=0`).
+Runs on **ZeroGPU** (NVIDIA RTX Pro 6000 Blackwell, `large` / 48 GB VRAM; 40
+min/day for Team org members). Each generation is capped at
+`PHANTOM_GRID_ZEROGPU_DURATION` seconds (default 90). No voice path
+(`PHANTOM_GRID_WITNESS_CHAT_TTS=0`).
 
 ## 🎥 Demo video
 

llm/zerogpu_backend.py

@@ -1,56 +1,60 @@
 """In-process transformers backend for HF Spaces ZeroGPU.
 
-Loaded only when PHANTOM_GRID_LLM_PROVIDER=zerogpu_transformers. Model weights
-live in CPU RAM and are moved to cuda inside the @spaces.GPU function, which is
-how ZeroGPU's per-call GPU acquisition expects PyTorch models to behave.
+Loaded only when PHANTOM_GRID_LLM_PROVIDER=zerogpu_transformers. Follows the
+canonical ZeroGPU pattern documented at
+https://huggingface.co/docs/hub/en/spaces-zerogpu :
+
+    "Models must be placed on `cuda` at the root module level. A PyTorch CUDA
+    emulation mode is enabled outside @spaces.GPU functions, allowing CUDA
+    operations without a real GPU. Inside @spaces.GPU, real CUDA is used."
+
+So we load directly to `cuda` here; ZeroGPU's emulation handles it at import,
+and the real device is attached only while the @spaces.GPU function runs.
 """
 from __future__ import annotations
 
-import json
 import os
-import re
-import threading
 from typing import Any
 
 _MODEL_ID = os.getenv("PHANTOM_GRID_ZEROGPU_MODEL_ID", "openbmb/MiniCPM4.1-8B")
 _DEFAULT_DURATION = int(os.getenv("PHANTOM_GRID_ZEROGPU_DURATION", "90"))
 
-_lock = threading.Lock()
-_model = None
-_tokenizer = None
-_loaded_to_cuda = False
-
-
-def _ensure_loaded() -> None:
-    """Load the model + tokenizer onto CPU. Called once at module init."""
-    global _model, _tokenizer
-    if _model is not None and _tokenizer is not None:
-        return
-    import torch
-    from transformers import AutoModelForCausalLM, AutoTokenizer
-
-    with _lock:
-        if _tokenizer is None:
-            _tokenizer = AutoTokenizer.from_pretrained(_MODEL_ID, trust_remote_code=True)
-        if _model is None:
-            _model = AutoModelForCausalLM.from_pretrained(
-                _MODEL_ID,
-                torch_dtype=torch.bfloat16,
-                trust_remote_code=True,
-                low_cpu_mem_usage=True,
-            )
-            _model.eval()
-
-
+# ZeroGPU's `spaces` package patches torch's CUDA at import time. We must import
+# it BEFORE torch so the emulation is in place when we then load to 'cuda'.
 try:
     import spaces  # type: ignore
-
     _spaces_gpu = spaces.GPU
-except ImportError:  # pragma: no cover — local dev without `spaces` installed
+    _ON_ZEROGPU = True
+except ImportError:  # local dev without the spaces package
     def _spaces_gpu(*_args, **_kwargs):
         def _wrap(fn):
             return fn
         return _wrap
+    _ON_ZEROGPU = False
+
+import torch  # noqa: E402  — must come after `import spaces`
+from transformers import AutoModelForCausalLM, AutoTokenizer  # noqa: E402
+
+_LOAD_ERROR: str | None = None
+_model = None
+_tokenizer = None
+
+try:
+    _tokenizer = AutoTokenizer.from_pretrained(_MODEL_ID, trust_remote_code=True)
+    # On ZeroGPU, `device_map="cuda"` lands on the emulated device at import
+    # and is 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"
+    _model = AutoModelForCausalLM.from_pretrained(
+        _MODEL_ID,
+        torch_dtype=torch.bfloat16,
+        trust_remote_code=True,
+        device_map=_device_map,
+        low_cpu_mem_usage=True,
+    )
+    _model.eval()
+except Exception as exc:  # pragma: no cover
+    _LOAD_ERROR = f"{exc.__class__.__name__}: {exc}"
 
 
 @_spaces_gpu(duration=_DEFAULT_DURATION)
@@ -60,26 +64,18 @@ def _generate_on_gpu(
     temperature: float,
     json_mode: bool,
 ) -> str:
-    """Runs on the H200 only when ZeroGPU has granted the slot."""
-    import torch
-
-    global _loaded_to_cuda
-    _ensure_loaded()
-    assert _model is not None and _tokenizer is not None
-
-    if not _loaded_to_cuda:
-        _model.to("cuda")
-        _loaded_to_cuda = True
+    if _model is None or _tokenizer is None:
+        raise RuntimeError(_LOAD_ERROR or "zerogpu backend not initialized")
 
     prompt = _tokenizer.apply_chat_template(
         messages,
         tokenize=False,
         add_generation_prompt=True,
     )
-    inputs = _tokenizer(prompt, return_tensors="pt").to("cuda")
+    inputs = _tokenizer(prompt, return_tensors="pt").to(_model.device)
 
     do_sample = temperature > 0
-    gen_kwargs = {
+    gen_kwargs: dict[str, Any] = {
         "max_new_tokens": max_new_tokens,
         "do_sample": do_sample,
         "pad_token_id": _tokenizer.eos_token_id,
@@ -93,14 +89,12 @@ def _generate_on_gpu(
 
     new_tokens = outputs[0][inputs["input_ids"].shape[1]:]
     text = _tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
-
     if json_mode:
         text = _extract_json(text)
     return text
 
 
 def _extract_json(text: str) -> str:
-    """Strip code fences and isolate the first JSON object."""
     cleaned = text.strip()
     if cleaned.startswith("```"):
         cleaned = cleaned.strip("`")
@@ -138,26 +132,18 @@ def chat_completion(
     max_tokens: int = 512,
     json_mode: bool = False,
 ) -> str:
-    """Synchronous, drop-in replacement for an OpenAI chat-completions POST."""
+    """Synchronous drop-in for an OpenAI chat-completions POST."""
     messages = _maybe_inject_no_think([dict(m) for m in messages])
     return _generate_on_gpu(messages, max_tokens, temperature, json_mode)
 
 
 def health() -> dict[str, Any]:
-    """Cheap probe — does NOT touch the GPU. Reflects whether the module imported."""
     return {
         "reachable": True,
-        "ready": _tokenizer is not None,
-        "detail": {"model_id": _MODEL_ID, "loaded_to_cuda": _loaded_to_cuda},
+        "ready": _model is not None and _LOAD_ERROR is None,
+        "detail": {
+            "model_id": _MODEL_ID,
+            "on_zerogpu": _ON_ZEROGPU,
+            "load_error": _LOAD_ERROR,
+        },
     }
-
-
-# Kick off the CPU-side load eagerly so the first user turn doesn't pay for it.
-# Wrapped in try/except so an import failure (e.g. during local dev without the
-# huge model cache) doesn't crash app startup — the health probe will report it.
-try:
-    _ensure_loaded()
-except Exception as exc:  # pragma: no cover
-    _LOAD_ERROR = f"{exc.__class__.__name__}: {exc}"
-else:
-    _LOAD_ERROR = None