Daankular commited on
Commit
30a271d
·
1 Parent(s): 6ea654b

Fix RMBG-2.0: use briaai/RMBG-2.0 (official) with standard load pattern

Browse files
Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -249,11 +249,11 @@ def load_triposg():
249
 
250
  try:
251
  from transformers import AutoModelForImageSegmentation
252
- # device_map loads weights directly to GPU, avoiding the meta tensor
253
- # intermediate that causes "Tensor.item() cannot be called on meta tensors"
254
  _rmbg_net = AutoModelForImageSegmentation.from_pretrained(
255
- "1038lab/RMBG-2.0", trust_remote_code=True, device_map=DEVICE
256
  )
 
 
257
  _rmbg_net.eval()
258
  _rmbg_version = "2.0"
259
  print("[load_triposg] TripoSG + RMBG-2.0 loaded.")
 
249
 
250
  try:
251
  from transformers import AutoModelForImageSegmentation
 
 
252
  _rmbg_net = AutoModelForImageSegmentation.from_pretrained(
253
+ "briaai/RMBG-2.0", trust_remote_code=True
254
  )
255
+ torch.set_float32_matmul_precision("high")
256
+ _rmbg_net.to(DEVICE)
257
  _rmbg_net.eval()
258
  _rmbg_version = "2.0"
259
  print("[load_triposg] TripoSG + RMBG-2.0 loaded.")