Trying to add CPU support.

app.py

@@ -17,10 +17,15 @@ from briarmbg import BriaRMBG
 from depth_anything_v2.dpt import DepthAnythingV2
 
 
-net = BriaRMBG.from_pretrained("briaai/RMBG-1.4")
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-net.to(device)
-net.eval()
+net_cpu = BriaRMBG.from_pretrained("briaai/RMBG-1.4")
+net_cpu.to('cpu')
+net_cpu.eval()
+
+net_gpu = None
+if torch.cuda.is_available():
+    net_gpu = BriaRMBG.from_pretrained("briaai/RMBG-1.4")
+    net_gpu.to('cuda')
+    net_gpu.eval()
 
 def resize_image(image):
     image = image.convert('RGB')
@@ -28,36 +33,45 @@ def resize_image(image):
     image = image.resize(model_input_size, Image.BILINEAR)
     return image
 
-@spaces.GPU(duration=6)
-def process_background(image):
-
-    # prepare input
-    orig_image = Image.fromarray(image)
-    w,h = orig_im_size = orig_image.size
-    image = resize_image(orig_image)
-    im_np = np.array(image)
-    im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2,0,1)
-    im_tensor = torch.unsqueeze(im_tensor,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])
-    if torch.cuda.is_available():
-        im_tensor=im_tensor.cuda()
-
-    #inference
-    result=net(im_tensor)
-    # post process
-    result = torch.squeeze(F.interpolate(result[0][0], size=(h,w), mode='bilinear') ,0)
-    ma = torch.max(result)
-    mi = torch.min(result)
-    result = (result-mi)/(ma-mi)
-    # image to pil
-    result_array = (result*255).cpu().data.numpy().astype(np.uint8)
+def _run_rmbg_on_image(image_np, net, device_str):
+    """Shared helper: run RMBG net on a numpy image and return a PIL RGBA with alpha mask."""
+    orig_image = Image.fromarray(image_np)
+    w, h = orig_image.size
+    img = resize_image(orig_image)
+    im_np = np.array(img)
+    im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0) / 255.0
+    im_tensor = normalize(im_tensor, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
+    if device_str == 'cuda':
+        im_tensor = im_tensor.cuda()
+    with torch.no_grad():
+        result = net(im_tensor)
+    result = torch.squeeze(F.interpolate(result[0][0], size=(h, w), mode='bilinear'), 0)
+    ma = torch.max(result); mi = torch.min(result)
+    result = (result - mi) / (ma - mi + 1e-8)
+    result_array = (result * 255).cpu().numpy().astype(np.uint8)
     pil_mask = Image.fromarray(np.squeeze(result_array))
-    # add the mask on the original image as alpha channel
     new_im = orig_image.copy()
     new_im.putalpha(pil_mask)
     return new_im
 
+@spaces.GPU(duration=6)
+def process_background_gpu(image):
+    if net_gpu is None:
+        raise RuntimeError("No GPU instance available")
+    return _run_rmbg_on_image(image, net_gpu, 'cuda')
+
+def process_background_cpu(image):
+    return _run_rmbg_on_image(image, net_cpu, 'cpu')
+
+# wrapper used by the UI: try GPU first, fall back to CPU on any exception
+def process_background(image):
+    try:
+        # attempt GPU call (this can raise if Zero-GPU is unavailable)
+        return process_background_gpu(image)
+    except Exception:
+        # fallback to CPU path
+        return process_background_cpu(image)
+
 
 css = """
 #img-display-container {