Update app_3.py

app_3.py

@@ -366,8 +366,8 @@ def encode_prompt_pair(positive_prompt, negative_prompt):
 
     return c, uc
 
-@spaces.GPU(duration=60)
-@torch.inference_mode()
+# @spaces.GPU(duration=60)
+# @torch.inference_mode()
 @spaces.GPU(duration=60)
 @torch.inference_mode()
 def infer(
@@ -467,29 +467,44 @@ def resize_without_crop(image, target_width, target_height):
 
 @spaces.GPU(duration=60)
 @torch.inference_mode()
-def run_rmbg(img, sigma=0.0):
-    # Convert RGBA to RGB if needed
-    if img.shape[-1] == 4:
-        # Use white background for alpha composition
-        alpha = img[..., 3:] / 255.0
-        rgb = img[..., :3]
-        white_bg = np.ones_like(rgb) * 255
-        img = (rgb * alpha + white_bg * (1 - alpha)).astype(np.uint8)
+# def run_rmbg(img, sigma=0.0):
+#     # Convert RGBA to RGB if needed
+#     if img.shape[-1] == 4:
+#         # Use white background for alpha composition
+#         alpha = img[..., 3:] / 255.0
+#         rgb = img[..., :3]
+#         white_bg = np.ones_like(rgb) * 255
+#         img = (rgb * alpha + white_bg * (1 - alpha)).astype(np.uint8)
     
-    H, W, C = img.shape
-    assert C == 3
-    k = (256.0 / float(H * W)) ** 0.5
-    feed = resize_without_crop(img, int(64 * round(W * k)), int(64 * round(H * k)))
-    feed = numpy2pytorch([feed]).to(device=device, dtype=torch.float32)
-    alpha = rmbg(feed)[0][0]
-    alpha = torch.nn.functional.interpolate(alpha, size=(H, W), mode="bilinear")
-    alpha = alpha.movedim(1, -1)[0]
-    alpha = alpha.detach().float().cpu().numpy().clip(0, 1)
+#     H, W, C = img.shape
+#     assert C == 3
+#     k = (256.0 / float(H * W)) ** 0.5
+#     feed = resize_without_crop(img, int(64 * round(W * k)), int(64 * round(H * k)))
+#     feed = numpy2pytorch([feed]).to(device=device, dtype=torch.float32)
+#     alpha = rmbg(feed)[0][0]
+#     alpha = torch.nn.functional.interpolate(alpha, size=(H, W), mode="bilinear")
+#     alpha = alpha.movedim(1, -1)[0]
+#     alpha = alpha.detach().float().cpu().numpy().clip(0, 1)
     
-    # Create RGBA image
-    rgba = np.dstack((img, alpha * 255)).astype(np.uint8)
-    result = 127 + (img.astype(np.float32) - 127 + sigma) * alpha
-    return result.clip(0, 255).astype(np.uint8), rgba
+#     # Create RGBA image
+#     rgba = np.dstack((img, alpha * 255)).astype(np.uint8)
+#     result = 127 + (img.astype(np.float32) - 127 + sigma) * alpha
+#     return result.clip(0, 255).astype(np.uint8), rgba
+
+@spaces.GPU
+@torch.inference_mode()
+def run_rmbg(image):
+    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()
+    pred = preds[0].squeeze()
+    pred_pil = transforms.ToPILImage()(pred)
+    mask = pred_pil.resize(image_size)
+    image.putalpha(mask)
+    return image
+  
 
 @spaces.GPU(duration=60)
 @torch.inference_mode()
@@ -615,9 +630,10 @@ def extract_foreground(image):
     if image is None:
         return None, gr.update(visible=True), gr.update(visible=True)
     logging.info(f"Input image shape: {image.shape}, dtype: {image.dtype}")
-    result, rgba = run_rmbg(image)
+    #result, rgba = run_rmbg(image)
+    result = run_rmbg(image)
     logging.info(f"Result shape: {result.shape}, dtype: {result.dtype}")
-    logging.info(f"RGBA shape: {rgba.shape}, dtype: {rgba.dtype}")
+    #logging.info(f"RGBA shape: {rgba.shape}, dtype: {rgba.dtype}")
     return result, gr.update(visible=True), gr.update(visible=True)
 
 def update_extracted_fg_height(selected_image: gr.SelectData):