Update inference_i2mv_sdxl.py

inference_i2mv_sdxl.py

@@ -103,17 +103,52 @@ def remove_bg(image: Image.Image, net, transform, device, mask: Image.Image = No
     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)[0].sigmoid().cpu()
-        #preds = net(input_images)[-1] if isinstance(net(input_images), list) else net(input_images)
-    pred = preds[0].squeeze()
-    pred_pil = transforms.ToPILImage()(pred)
-    mask = pred_pil.resize(image_size)
-    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)[0].sigmoid().cpu()
+#         #preds = net(input_images)[-1] if isinstance(net(input_images), list) else net(input_images)
+#     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):
+#     """
+#     Applies a pre-existing mask to an image to make the background transparent.
+#     Args:
+#         image (PIL.Image.Image): The input image.
+#         net: Pre-trained neural network (not used but kept for compatibility).
+#         transform: Image transformation object (not used but kept for compatibility).
+#         device: Device used for inference (not used but kept for compatibility).
+#     Returns:
+#         PIL.Image.Image: The modified image with transparent background.
+#     """
+#     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)
+    
+#     # Resize the mask to match the original image size
+#     mask = pred_pil.resize(image_size, Image.LANCZOS)
+    
+#     # Create a new image with the same size and mode as the original
+#     output_image = Image.new("RGBA", image_size)
+    
+#     # Apply the mask to the original image
+#     image.putalpha(mask)
+    
+#     # Composite the original image with the mask
+#     output_image.paste(image, (0, 0), image)
+    
+#     return output_image