rembg->carvekit gpu

app.py

@@ -74,7 +74,7 @@ from pipeline_zero1to3 import Zero1to3StableDiffusionPipeline
 from segment_anything import sam_model_registry, SamPredictor
 
 import rembg
-# from carvekit.api.high import HiInterface
+from carvekit.api.high import HiInterface
 
 
 pretrained_model_name_or_path = "kxic/EscherNet_demo"
@@ -130,25 +130,31 @@ def sam_init():
     predictor = SamPredictor(sam)
     return predictor
 
-# @spaces.GPU
-# def create_carvekit_interface():
-#     # Check doc strings for more information
-#     interface = HiInterface(object_type="object",  # Can be "object" or "hairs-like".
-#                             batch_size_seg=6,
-#                             batch_size_matting=1,
-#                             device=device,
-#                             seg_mask_size=640,  # Use 640 for Tracer B7 and 320 for U2Net
-#                             matting_mask_size=2048,
-#                             trimap_prob_threshold=231,
-#                             trimap_dilation=30,
-#                             trimap_erosion_iters=5,
-#                             fp16=True)
-#
-#     return interface
-
-
-rembg_session = rembg.new_session()
-# rembg_session = create_carvekit_interface()
+def create_carvekit_interface():
+    # Check doc strings for more information
+    interface = HiInterface(object_type="object",  # Can be "object" or "hairs-like".
+                            batch_size_seg=6,
+                            batch_size_matting=1,
+                            device="cpu",
+                            seg_mask_size=640,  # Use 640 for Tracer B7 and 320 for U2Net
+                            matting_mask_size=2048,
+                            trimap_prob_threshold=231,
+                            trimap_dilation=30,
+                            trimap_erosion_iters=5,
+                            fp16=False)
+
+    return interface
+
+
+# rembg_session = rembg.new_session()
+rembg_session = create_carvekit_interface()
+rembg_session.u2net = rembg_session.u2net.to(device)
+rembg_session.fba = rembg_session.fba.to(device)
+rembg_session.fba.device = device
+rembg_session.device = device
+rembg_session.u2net.device = device
+# rembg_session.postprocessing_pipeline = rembg_session.postprocessing_pipeline.to(device)
+# rembg_session.postprocessing_pipeline.device = device
 
 predictor = sam_init()
 

dust3r/utils/image.py

@@ -119,9 +119,9 @@ def load_images(folder_or_list, size, square_ok=False, verbose=True, do_remove_b
         # remove background if needed
         if do_remove_background:
             # use rembg
-            image_nobg = remove(img, alpha_matting=True, session=rembg_session)
+            # image_nobg = remove(img, alpha_matting=True, session=rembg_session)
             # use carvekit
-            # image_nobg = rembg_session([img])[0]
+            image_nobg = rembg_session([img])[0]
             arr = np.asarray(image_nobg)[:, :, -1]
             x_nonzero = np.nonzero(arr.sum(axis=0))
             y_nonzero = np.nonzero(arr.sum(axis=1))