space: load the released nf4 Warden from the hub (1GB Space LFS cap rules out in-repo weights)

space/app.py

@@ -57,15 +57,20 @@ STATIC = Path(__file__).parent / "static"
 
 # ------------------------------------------------------------ the Warden brain
 
-# Prefer the Warden shipped inside the Space repo (pre-quantized nf4, no
-# download at boot); fall back to the hub for self-hosters who clone without
-# the LFS weights.
+# Model source, in preference order:
+#   1. WARDEN_MODEL env (+ optional WARDEN_SUBFOLDER)
+#   2. weights shipped in the repo at model/ (only possible with persistent
+#      storage — Space repos cap LFS at 1GB, so this dir is normally absent)
+#   3. the released Warden, pre-quantized nf4 (~18GB, HF-internal download —
+#      this is what kills the boot-time 63GB download + quantize wait)
 _SHIPPED = REPO_ROOT / "model"
-MODEL_ID = os.environ.get("WARDEN_MODEL") or (
-    str(_SHIPPED)
-    if (_SHIPPED / "config.json").exists()
-    else "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16"
-)
+if os.environ.get("WARDEN_MODEL"):
+    MODEL_ID = os.environ["WARDEN_MODEL"]
+    SUBFOLDER = os.environ.get("WARDEN_SUBFOLDER", "")
+elif any(_SHIPPED.glob("*.safetensors")):
+    MODEL_ID, SUBFOLDER = str(_SHIPPED), ""
+else:
+    MODEL_ID, SUBFOLDER = "IMJONEZZ/warden-nemotron-3-nano-30b", "bnb-4bit"
 INTERNAL_KEY = os.environ.get("SCRYPT_INTERNAL_KEY") or secrets.token_hex(16)
 
 tok = None
@@ -128,21 +133,26 @@ if spaces is not None:
         _ensure_mamba_kernels()
         import torch
         from transformers import (
+            AutoConfig,
             AutoModelForCausalLM,
             AutoTokenizer,
             BitsAndBytesConfig,
         )
 
-        tok = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+        tok = AutoTokenizer.from_pretrained(
+            MODEL_ID, subfolder=SUBFOLDER, trust_remote_code=True
+        )
         # 4-bit fits the `large` (48GB) slice and matches the Q4 GGUF the
-        # local game ships — quality parity, not a downgrade. The shipped
-        # model/ dir is already nf4 (config carries quantization_config);
-        # only quantize on the fly when loading raw BF16 from the hub.
-        _cfg = Path(MODEL_ID) / "config.json"
-        _already_quant = _cfg.exists() and "quantization_config" in _cfg.read_text()
+        # local game ships — quality parity, not a downgrade. The released
+        # bnb-4bit weights are already nf4 (config carries
+        # quantization_config); only quantize on the fly when pointed at
+        # raw BF16 weights.
+        _cfg = AutoConfig.from_pretrained(
+            MODEL_ID, subfolder=SUBFOLDER, trust_remote_code=True
+        )
         _quant_kwargs = (
             {}
-            if _already_quant
+            if getattr(_cfg, "quantization_config", None)
             else {
                 "quantization_config": BitsAndBytesConfig(
                     load_in_4bit=True,
@@ -153,6 +163,7 @@ if spaces is not None:
         )
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_ID,
+            subfolder=SUBFOLDER,
             trust_remote_code=True,
             device_map="cuda",
             **_quant_kwargs,
@@ -221,7 +232,7 @@ def status() -> dict:
         "warden_ready": WARDEN_READY,
         "warden_error": WARDEN_ERR,
         "mamba": MAMBA_DIAG,
-        "model": MODEL_ID,
+        "model": MODEL_ID + (f"/{SUBFOLDER}" if SUBFOLDER else ""),
         "zerogpu": spaces is not None,
     }