Update app.py

app.py

@@ -7,11 +7,10 @@ from torchvision import transforms
 from typing import Union, Tuple
 from PIL import Image
 
-torch.set_float32_matmul_precision(["high", "highest"][0])
-
 birefnet = AutoModelForImageSegmentation.from_pretrained(
-    "ZhengPeng7/BiRefNet", trust_remote_code=True
+    "ZhengPeng7/BiRefNet", low_cpu_mem_usage=False,  trust_remote_code=True, torch_dtype=torch.float32, device_map=None
 )
+birefnet = birefnet.eval()
 #birefnet.to("cuda")
 
 transform_image = transforms.Compose(
@@ -59,10 +58,9 @@ def process(image: Image.Image) -> Image.Image:
         PIL.Image: The image with the background removed, using the segmentation mask as transparency.
     """
     image_size = image.size
-    input_images = transform_image(image).unsqueeze(0)#.to("cuda")
-    # Prediction
-    with torch.no_grad():
-        preds = birefnet(input_images)[-1].sigmoid().cpu()
+    input_images = transform_image(image).unsqueeze(0)
+    with torch.inference_mode():
+        preds = birefnet(input_images)[-1].sigmoid().detach().cpu()
     pred = preds[0].squeeze()
     pred_pil = transforms.ToPILImage()(pred)
     mask = pred_pil.resize(image_size)