Spaces:

thelabel
/

background-remover

Sleeping

user-agent commited on 5 days ago

Commit

8239fe8

verified ·

1 Parent(s): 98c6c6d

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -13,7 +13,8 @@ torch.set_float32_matmul_precision("high")
 birefnet = AutoModelForImageSegmentation.from_pretrained(
     "ZhengPeng7/BiRefNet", trust_remote_code=True
 )
-birefnet.to("cuda")
 # Image transformation
 transform_image = transforms.Compose([

 birefnet = AutoModelForImageSegmentation.from_pretrained(
     "ZhengPeng7/BiRefNet", trust_remote_code=True
 )
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+birefnet.to(device)
 # Image transformation
 transform_image = transforms.Compose([