Update inference_i2mv_sdxl.py

inference_i2mv_sdxl.py

@@ -71,16 +71,44 @@ def prepare_pipeline(
     return pipe
 
 
-def remove_bg(image, net, transform, device):
+# 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.
+
+    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).
+        mask (PIL.Image.Image, optional): The mask to use. Should be the same size
+                                 as the input image, with values between 0 and 255 (or 0-1).
+                                 If None, will return image with no changes.
+
+    Returns:
+        PIL.Image.Image: The modified image with transparent background.
+    """
+    if mask is None:
+      return image
+    
     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)
+    if mask.size != image_size:
+        mask = mask.resize(image_size) # Resizing the mask if it is not the same size as image
+    
     image.putalpha(mask)
     return image
+    
 
 
 def preprocess_image(image: Image.Image, height, width):
@@ -150,7 +178,8 @@ def run_pipeline(
     # Prepare image
     reference_image = Image.open(image) if isinstance(image, str) else image
     if remove_bg_fn is not None:
-        reference_image = remove_bg_fn(reference_image)
+        # reference_image = remove_bg_fn(reference_image)
+        
         reference_image = preprocess_image(reference_image, height, width)
     elif reference_image.mode == "RGBA":
         reference_image = preprocess_image(reference_image, height, width)