Patch image_process.py rmbg_net None guard; fix RMBG-2.0 loading with device_map

app.py

@@ -191,6 +191,26 @@ def load_triposg():
     if str(triposg_src) not in sys.path:
         sys.path.insert(0, str(triposg_src))
 
+    # Patch image_process.py: guard rmbg_net=None in load_image.
+    # TripoSG calls rmbg(rgb_image_resized) unconditionally when alpha is None,
+    # with no check for rmbg_net being None. Fallback: all-white alpha (full foreground).
+    _ip_path = triposg_src / "scripts" / "image_process.py"
+    if _ip_path.exists():
+        _ip_text = _ip_path.read_text()
+        if "rmbg_net_none_guard" not in _ip_text:
+            _ip_text = _ip_text.replace(
+                "        # seg from rmbg\n        alpha_gpu_rmbg = rmbg(rgb_image_resized)",
+                "        # seg from rmbg\n"
+                "        if rmbg_net is None:  # rmbg_net_none_guard\n"
+                "            alpha_gpu_rmbg = torch.ones(\n"
+                "                1, rgb_image_resized.shape[1], rgb_image_resized.shape[2],\n"
+                "                device=rgb_image_resized.device)\n"
+                "        else:\n"
+                "            alpha_gpu_rmbg = rmbg(rgb_image_resized)",
+            )
+            _ip_path.write_text(_ip_text)
+            print("[load_triposg] Patched image_process.py: rmbg_net None guard")
+
     # Safety net: patch inference_utils.py to make diso import optional.
     # Even if diso compiled with submodules, guard against any residual link errors.
     _iu_path = triposg_src / "triposg" / "inference_utils.py"
@@ -229,9 +249,11 @@ def load_triposg():
 
     try:
         from transformers import AutoModelForImageSegmentation
+        # device_map loads weights directly to GPU, avoiding the meta tensor
+        # intermediate that causes "Tensor.item() cannot be called on meta tensors"
         _rmbg_net = AutoModelForImageSegmentation.from_pretrained(
-            "1038lab/RMBG-2.0", trust_remote_code=True, low_cpu_mem_usage=False
-        ).to(DEVICE)
+            "1038lab/RMBG-2.0", trust_remote_code=True, device_map=DEVICE
+        )
         _rmbg_net.eval()
         _rmbg_version = "2.0"
         print("[load_triposg] TripoSG + RMBG-2.0 loaded.")