Update app.py

app.py

@@ -14,7 +14,7 @@ torch.set_float32_matmul_precision(["high", "highest"][0])
 birefnet = AutoModelForImageSegmentation.from_pretrained(
     "ZhengPeng7/BiRefNet", trust_remote_code=True
 )
-birefnet.to("cuda")
+birefnet.to("cpu")
 transform_image = transforms.Compose(
     [
         transforms.Resize((1024, 1024)),
@@ -27,7 +27,7 @@ def fn(image):
     im = load_img(image, output_type="pil")
     im = im.convert("RGB")
     image_size = im.size
-    input_images = transform_image(im).unsqueeze(0).to("cuda")
+    input_images = transform_image(im).unsqueeze(0).to("cpu")
     
     with torch.no_grad():
         preds = birefnet(input_images)[-1].sigmoid().cpu()
@@ -48,7 +48,7 @@ def fn_url(url):
     im = load_img(url, output_type="pil")
     im = im.convert("RGB")
     image_size = im.size
-    input_images = transform_image(im).unsqueeze(0).to("cuda")
+    input_images = transform_image(im).unsqueeze(0).to("cpu")
     
     with torch.no_grad():
         preds = birefnet(input_images)[-1].sigmoid().cpu()
@@ -71,7 +71,7 @@ def batch_fn(images):
         im = load_img(image_path, output_type="pil")
         im = im.convert("RGB")
         image_size = im.size
-        input_images = transform_image(im).unsqueeze(0).to("cuda")
+        input_images = transform_image(im).unsqueeze(0).to("cpu")
         
         with torch.no_grad():
             preds = birefnet(input_images)[-1].sigmoid().cpu()