From bc99c01df5612b083be3211ff2e3cee34cd608eb Mon Sep 17 00:00:00 2001
From: AI Lab <129358391+1038lab@users.noreply.github.com>
Date: Sun, 29 Dec 2024 22:01:05 -0800
Subject: [PATCH] Add files via upload

---
 AILab_RMBG.py    | 13 ++++++++-----
 AILab_Segment.py |  5 +++++
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/AILab_RMBG.py b/AILab_RMBG.py
index 2f77323..8fcdaf9 100644
--- a/AILab_RMBG.py
+++ b/AILab_RMBG.py
@@ -376,7 +376,7 @@ class RMBG:
                 # Get mask from specific model
                 mask = model_instance.process_image(img, model, params)
                 
-                # check if mask is a list
+                # Ensure mask is in the correct format
                 if isinstance(mask, list):
                     masks = [m.convert("L") for m in mask if isinstance(m, Image.Image)]
                     mask = masks[0] if masks else None
@@ -402,19 +402,22 @@ class RMBG:
                 
                 if params["invert_output"]:
                     mask = Image.fromarray(255 - np.array(mask))
-                
+
                 # Create final image
                 orig_image = tensor2pil(img)
                 orig_rgba = orig_image.convert("RGBA")
                 r, g, b, _ = orig_rgba.split()
                 foreground = Image.merge('RGBA', (r, g, b, mask))
-                
+
                 if params["background"] != "Alpha":
                     bg_color = bg_colors[params["background"]]
                     bg_image = Image.new('RGBA', orig_image.size, (*bg_color, 255))
                     composite_image = Image.alpha_composite(bg_image, foreground)
-                    processed_images.append(pil2tensor(composite_image))
+                    
+                    # Convert to RGB if background is not Alpha
+                    processed_images.append(pil2tensor(composite_image.convert("RGB")))
                 else:
+                    # Keep as RGBA if background is Alpha
                     processed_images.append(pil2tensor(foreground))
                 
                 processed_masks.append(pil2tensor(mask))
@@ -430,5 +433,5 @@ NODE_CLASS_MAPPINGS = {
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
-    "RMBG": "RMBG (Background Remover)"
+    "RMBG": "Remove Background (RMBG)"
 } 
\ No newline at end of file
diff --git a/AILab_Segment.py b/AILab_Segment.py
index ec415a2..413c8fc 100644
--- a/AILab_Segment.py
+++ b/AILab_Segment.py
@@ -205,6 +205,11 @@ class Segment:
         
         result_image = apply_background_color(image, mask_image, background_color)
         
+        if background_color != "Alpha":
+            result_image = result_image.convert("RGB")
+        else:
+            result_image = result_image.convert("RGBA")
+
         print(f'Successfully created segment for: "{prompt}"')
         return (pil2tensor(result_image), image2mask(mask_image))