fix(spaces): Hdri torch.load after ConvNeXt; default off CPU preload with spaces

- hdri_encoder: torch.load(..., weights_only=False, map_location=cpu) so Stateless
GPU WeightsUnpickler does not run after torchvision may have touched CUDA
- app: default NEAR_MODEL_CPU_PRELOAD_AT_START=0 when spaces is installed
- test: assert Hdri load_weights uses full unpickle on CPU

Made-with: Cursor

app.py

@@ -253,10 +253,18 @@ def _truthy_env(name: str, default: str) -> bool:
     return v in ("1", "true", "yes", "on")
 
 
-_CPU_PRELOAD_AT_START = _truthy_env("NEAR_MODEL_CPU_PRELOAD_AT_START", "1")
+# Background CPU preload runs in the Gradio host process. HF Stateless / ZeroGPU forbids
+# CUDA init there; torchvision + torch.load (weights_only) interactions can still fail
+# during NeAR build. Default to off when `spaces` is present (typical Space deploy); use
+# NEAR_MODEL_CPU_PRELOAD_AT_START=1 to force preload locally or on dedicated GPU VMs.
+_CPU_PRELOAD_DEFAULT = "0" if spaces is not None else "1"
+_CPU_PRELOAD_AT_START = _truthy_env(
+    "NEAR_MODEL_CPU_PRELOAD_AT_START",
+    _CPU_PRELOAD_DEFAULT,
+)
 print(
     f"[NeAR] NEAR_MODEL_CPU_PRELOAD_AT_START={'1' if _CPU_PRELOAD_AT_START else '0'} "
-    "(Hunyuan + NeAR weights on CPU at process start; GPU callbacks only .to(cuda) + infer).",
+    f"(default {_CPU_PRELOAD_DEFAULT!r} when spaces={'set' if spaces is not None else 'absent'}).",
     flush=True,
 )
 

tests/test_hdri_encoder_torch_load.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+import unittest
+from pathlib import Path
+
+
+HDRI_PATH = Path(__file__).resolve().parents[1] / "trellis" / "models" / "structured_latent_vae" / "hdri_encoder.py"
+
+
+class HdriEncoderTorchLoadTests(unittest.TestCase):
+    def test_load_weights_uses_full_unpickle_on_cpu(self) -> None:
+        text = HDRI_PATH.read_text(encoding="utf-8")
+        self.assertIn("weights_only=False", text)
+        self.assertIn("torch.device(\"cpu\")", text)
+
+
+if __name__ == "__main__":
+    unittest.main()

trellis/models/structured_latent_vae/hdri_encoder.py

@@ -374,7 +374,15 @@ class Hdri_Encoder(nn.Module):
 
     def load_weights(self, pretrained_path):
         if pretrained_path is not None:
-            checkpoint = torch.load(pretrained_path, map_location='cpu')
+            # ConvNeXt ImageNet weights are loaded above; that can initialize CUDA in the
+            # process. On Hugging Face Stateless / ZeroGPU, the next torch.load with
+            # weights_only=True (PyTorch 2.6+ default) then fails inside WeightsUnpickler.
+            # This checkpoint is from the NeAR bundle (trusted); full unpickle on CPU only.
+            checkpoint = torch.load(
+                pretrained_path,
+                map_location=torch.device("cpu"),
+                weights_only=False,
+            )
             self.load_state_dict(checkpoint)
 
     def forward(self, context):