diff --git a/AILab_RMBG.py b/AILab_RMBG.py index fcfd71e..2f77323 100644 --- a/AILab_RMBG.py +++ b/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"])) diff --git a/AILab_Segment.py b/AILab_Segment.py index d237d1c..ec415a2 100644 --- a/AILab_Segment.py +++ b/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. diff --git a/update.md b/update.md index 84076ee..3032f96 100644 --- a/update.md +++ b/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"