briaai/RMBG-1.4 always run on cpu

This commit is contained in:
Qing 2024-02-10 13:19:43 +08:00
parent 893fd11b98
commit 1cd36cbe60

View File

@ -490,9 +490,6 @@ def briarmbg_process(bgr_np_image, session, only_mask=False):
im_tensor = torch.unsqueeze(im_tensor, 0) im_tensor = torch.unsqueeze(im_tensor, 0)
im_tensor = torch.divide(im_tensor, 255.0) im_tensor = torch.divide(im_tensor, 255.0)
im_tensor = normalize(im_tensor, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0]) im_tensor = normalize(im_tensor, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
if torch.cuda.is_available():
im_tensor = im_tensor.cuda()
# inference # inference
result = session(im_tensor) result = session(im_tensor)
# post process # post process