Update inference_i2mv_sdxl.py

inference_i2mv_sdxl.py

@@ -70,19 +70,6 @@ def prepare_pipeline(
 
     return pipe
 
-
-# def remove_bg(image, net, transform, device):
-#     image_size = image.size
-#     input_images = transform(image).unsqueeze(0).to(device)
-#     with torch.no_grad():
-#         preds = net(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
-
-
 def remove_bg(image: Image.Image, net, transform, device, mask: Image.Image = None):
     """
     Applies a pre-existing mask to an image to make the background transparent.
@@ -108,21 +95,35 @@ def remove_bg(image: Image.Image, net, transform, device, mask: Image.Image = No
     
     image.putalpha(mask)
     return image
-    
+
+
+def remove_bg(image, net, transform, device):
+    image_size = image.size
+    input_images = transform(image).unsqueeze(0).to(device)
+    with torch.no_grad():
+        preds = net(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
+
+
+
 
 
 def preprocess_image(image: Image.Image, height, width):
     
-    # alpha = image[..., 3] > 0
+    alpha = image[..., 3] > 0
     # alpha = image
 
-    if image.mode in ("RGBA", "LA"):
-        image = np.array(image)
-        alpha = image[..., 3]  # Extract the alpha channel
-    elif image.mode in ("RGB"):
-        image = np.array(image)
+    #if image.mode in ("RGBA", "LA"):
+    #    image = np.array(image)
+    #    alpha = image[..., 3]  # Extract the alpha channel
+    #elif image.mode in ("RGB"):
+    #    image = np.array(image)
         # Create default alpha for non-alpha images
-        alpha = np.ones(image[..., 0].shape, dtype=np.uint8) * 255 # Create
+    #    alpha = np.ones(image[..., 0].shape, dtype=np.uint8) * 255 # Create
     H, W = alpha.shape
     # get the bounding box of alpha
     y, x = np.where(alpha)