Add files via upload

AILab_RMBG.py

@@ -1,4 +1,4 @@
-# ComfyUI-RMBG v1.3.0
+# ComfyUI-RMBG v1.3.1
 # This custom node for ComfyUI provides functionality for background removal using various models,
 # including RMBG-2.0, INSPYRENET, and BEN. It leverages deep learning techniques
 # to process images and generate masks for background removal.
@@ -304,21 +304,6 @@ class BENModel(BaseModelLoader):
             handle_model_error(f"Error in BEN processing: {str(e)}")
 
 class RMBG:
-    """
-    RMBG Node: Advanced Background Removal Suite
-    
-    This node provides professional background removal capabilities using three state-of-the-art models:
-    - RMBG-2.0: Latest model with excellent performance on complex backgrounds
-    - INSPYRENET: Specialized for human portrait segmentation
-    - BEN: Versatile model with good balance of speed and accuracy
-    
-    Features:
-    - Batch processing support
-    - Multiple background options
-    - Advanced mask refinement
-    - High-quality edge preservation
-    - Memory-efficient processing
-    """
     def __init__(self):
         self.models = {
             "RMBG-2.0": RMBGModel(),
@@ -391,6 +376,13 @@ class RMBG:
                 # Get mask from specific model
                 mask = model_instance.process_image(img, model, params)
                 
+                # check if mask is a list
+                if isinstance(mask, list):
+                    masks = [m.convert("L") for m in mask if isinstance(m, Image.Image)]
+                    mask = masks[0] if masks else None
+                elif isinstance(mask, Image.Image):
+                    mask = mask.convert("L")
+
                 # Post-process mask
                 mask_tensor = pil2tensor(mask)
                 mask_tensor = mask_tensor * (1 + (1 - params["sensitivity"]))

AILab_Segment.py

@@ -1,4 +1,4 @@
-# ComfyUI-RMBG v1.3.0
+# ComfyUI-RMBG v1.3.1
 # This custom node for ComfyUI provides functionality for background removal using various models,
 # including RMBG-2.0, INSPYRENET, and BEN. It leverages deep learning techniques
 # to process images and generate masks for background removal.

update.md

@@ -1,5 +1,11 @@
 # ComfyUI-RMBG Update Log
 
+## v1.3.1 (2024/12/25)
+
+### Bug Fixes
+- Fixed an issue with mask processing when the model returns a list of masks.
+- Improved handling of image formats to prevent processing errors.
+
 ## v1.3.0 (2024/12/23)
 
 ### New Segment (RMBG) Node
@@ -8,8 +14,6 @@
   - Support for multiple object detection and segmentation
   - Perfect for precise object extraction and recognition tasks
 
-![rmbg_v1 3 0](https://github.com/user-attachments/assets/0bff9a4d-7e22-403e-9ffb-e9defb335d9f)
-
 ### Supported Models
 - SAM (Segment Anything Model)
   - sam_vit_h: 2.56GB - Highest accuracy
@@ -30,7 +34,6 @@
   - Smart GPU memory handling
 
 ### Usage Examples
-![rmbg v1.3.0](https://github.com/user-attachments/assets/7607546e-ffcb-45e2-ab90-83267292757e)
 1. Tag-Style Prompts
    - Single object: "cat"
    - Multiple objects: "cat, dog, person"