diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..333d52a --- /dev/null +++ b/__init__.py @@ -0,0 +1,16 @@ +from .background_erase_network import ( + BackgroundEraseNetwork, + NODE_CLASS_MAPPINGS as BACKGROUND_ERASE_NETWORK_MAPPINGS, + NODE_DISPLAY_NAME_MAPPINGS as BACKGROUND_ERASE_NETWORK_DISPLAY_MAPPINGS +) + +# Combined NODE_CLASS_MAPPINGS dictionary +NODE_CLASS_MAPPINGS = { + **BACKGROUND_ERASE_NETWORK_MAPPINGS + +} + +# Combined NODE_DISPLAY_NAME_MAPPINGS dictionary +NODE_DISPLAY_NAME_MAPPINGS = { + **BACKGROUND_ERASE_NETWORK_DISPLAY_MAPPINGS +} diff --git a/background_erase_network.py b/background_erase_network.py new file mode 100644 index 0000000..0823e8e --- /dev/null +++ b/background_erase_network.py @@ -0,0 +1,69 @@ +import sys +import os +import torch +from PIL import Image +from torchvision import transforms + +# Add the directory of the current file to sys.path +script_directory = os.path.dirname(os.path.realpath(__file__)) +sys.path.append(script_directory) + +# Now import BEN_Base from model.py +from model import BEN_Base + +class BackgroundEraseNetwork: + def __init__(self): + self.model = BEN_Base() # Initialize BEN_Base model + self.model.loadcheckpoints("./ComfyUI/custom_nodes/ComfyUI-BEN/BEN_Base.pth") # Load the model weights + + # Define transformations for converting between PIL and Tensor + self.to_pil = transforms.ToPILImage() + self.to_tensor = transforms.ToTensor() + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "input_image": ("IMAGE",), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "process_image" + CATEGORY = "BEN" + + def process_image(self, input_image): + # Handle the input tensor format from ComfyUI + if isinstance(input_image, torch.Tensor): + if input_image.dim() == 4: + input_image = input_image[0] + + if input_image.dim() == 3: + input_image = input_image.permute(2, 0, 1) + + input_image = self.to_pil(input_image) + + # Ensure the image is in RGBA mode + if input_image.mode != 'RGBA': + input_image = input_image.convert("RGBA") + + # Run inference to get the foreground image + _, foreground = self.model.inference(input_image) + + # Convert the foreground to tensor + foreground_tensor = self.to_tensor(foreground) + + # Convert to ComfyUI format [B, H, W, C] + foreground_tensor = foreground_tensor.permute(1, 2, 0).unsqueeze(0) + + return (foreground_tensor,) + +# Export mappings for ComfyUI +NODE_CLASS_MAPPINGS = { + "BackgroundEraseNetwork": BackgroundEraseNetwork +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "BackgroundEraseNetwork": "Background Erase Network" +} \ No newline at end of file diff --git a/ben.png b/ben.png new file mode 100644 index 0000000..2ce4085 Binary files /dev/null and b/ben.png differ diff --git a/readme.md b/readme.md new file mode 100644 index 0000000..9cfe279 --- /dev/null +++ b/readme.md @@ -0,0 +1,51 @@ +# ComfyUI BEN (Background Erase Network) + +**** + + +Remove backgrounds from images with [BEN](https://huggingface.co/PramaLLC/BEN) in [ComfyUI](https://github.com/comfyanonymous/ComfyUI) + +## Installation + +``` +git clone https://github.com/DoctorDiffusion/ComfyUI-BEN.git +``` +``` +cd ComfyUI-BEN +``` +``` +pip install -r requirements.txt +``` +Go to [BEN Huggingface page](https://huggingface.co/PramaLLC/BEN/tree/main) and download BEN_Base.pth +``` +``` +cd ComfyUI-BEN +``` +The folder should look like this: +└── ComfyUI/custom-nodes/ComfyUI-BEN/ + ├── __init__.py + ├── background_erase_network.py + ├── ben.png + ├── BEN_Base.pth + ├── model.py + ├── readme.txt + └── requirements.txt +``` +## Nodes + +### Background Erase Network + +Outputs image of subject with alpha layer, combine with Image to Mask Node to get alpha mask as well. + +## Credits + +- [comfyanonymous/ComfyUI](https://github.com/comfyanonymous/ComfyUI) + +- [PramaLLC/BEN](https://huggingface.co/PramaLLC/BEN) + +⭐ If you like the project, please give it a star! ⭐ + +## License + +While this is under an Apache 2.0 License, [PramaLLC](https://huggingface.co/PramaLLC/BEN) has restricted the use of their model for for commercial purposes. Reach out to them directly for more information. + diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..1531126 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,8 @@ +huggingface_hub +Image +numpy +torch +einops +Pillow +timm +torchvision \ No newline at end of file