From 0816df8277574658251ca343efbd5936fd179348 Mon Sep 17 00:00:00 2001 From: AI Lab <129358391+1038lab@users.noreply.github.com> Date: Fri, 7 Feb 2025 08:44:34 +0800 Subject: [PATCH] Add files via upload --- AILab_RMBG.py | 111 +++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 106 insertions(+), 5 deletions(-) diff --git a/AILab_RMBG.py b/AILab_RMBG.py index 4f1df92..055ee48 100644 --- a/AILab_RMBG.py +++ b/AILab_RMBG.py @@ -1,6 +1,6 @@ -# ComfyUI-RMBG v1.7.0 +# ComfyUI-RMBG v1.8.0 # This custom node for ComfyUI provides functionality for background removal using various models, -# including RMBG-2.0, INSPYRENET, BEN and BEN2. It leverages deep learning techniques +# including RMBG-2.0, INSPYRENET, BEN, BEN2 and BIREFNET-HR. It leverages deep learning techniques # to process images and generate masks for background removal. # # Models License Notice: @@ -8,7 +8,8 @@ # - INSPYRENET: MIT License (https://github.com/plemeri/InSPyReNet) # - BEN: Apache-2.0 License (https://huggingface.co/PramaLLC/BEN) # - BEN2: Apache-2.0 License (https://huggingface.co/PramaLLC/BEN2) -# +# - BIREFNET-HR: Apache-2.0 License (https://huggingface.co/ZhengPeng7/BiRefNet_HR) +# # This integration script follows GPL-3.0 License. # When using or modifying this code, please respect both the original model licenses # and this integration's license terms. @@ -39,7 +40,7 @@ folder_paths.add_model_folder_path("rmbg", os.path.join(folder_paths.models_dir, AVAILABLE_MODELS = { "RMBG-2.0": { "type": "rmbg", - "repo_id": "briaai/RMBG-2.0", + "repo_id": "1038lab/RMBG-2.0", "files": { "config.json": "config.json", "model.safetensors": "model.safetensors", @@ -73,6 +74,17 @@ AVAILABLE_MODELS = { "BEN2.py": "BEN2.py" }, "cache_dir": "BEN2" + }, + "BIREFNET-HR": { + "type": "birefnet", + "repo_id": "1038lab/BiRefNet_HR", + "files": { + "birefnet.py": "birefnet.py", + "BiRefNet_config.py": "BiRefNet_config.py", + "model.safetensors": "model.safetensors", + "config.json": "config.json" + }, + "cache_dir": "BIREFNET-HR" } } @@ -389,13 +401,102 @@ class BEN2Model(BaseModelLoader): except Exception as e: handle_model_error(f"Error in BEN2 processing: {str(e)}") +class BiRefNetModel(BaseModelLoader): + def __init__(self): + super().__init__() + + def load_model(self, model_name): + if self.current_model_version != model_name: + self.clear_model() + + cache_dir = self.get_cache_dir(model_name) + model_path = os.path.join(cache_dir, "birefnet.py") + config_path = os.path.join(cache_dir, "BiRefNet_config.py") + weights_path = os.path.join(cache_dir, "model.safetensors") + + try: + # Fix relative imports in model file + with open(model_path, 'r', encoding='utf-8') as f: + model_content = f.read() + model_content = model_content.replace("from .BiRefNet_config", "from BiRefNet_config") + with open(model_path, 'w', encoding='utf-8') as f: + f.write(model_content) + + # Load config and model dynamically + spec = importlib.util.spec_from_file_location("BiRefNet_config", config_path) + config_module = importlib.util.module_from_spec(spec) + sys.modules["BiRefNet_config"] = config_module + spec.loader.exec_module(config_module) + + spec = importlib.util.spec_from_file_location("birefnet", model_path) + model_module = importlib.util.module_from_spec(spec) + sys.modules["birefnet"] = model_module + spec.loader.exec_module(model_module) + + # Initialize model + self.model = model_module.BiRefNet(config_module.BiRefNetConfig()) + + # Load weights using safetensors + from safetensors.torch import load_file + state_dict = load_file(weights_path) + self.model.load_state_dict(state_dict) + + self.model.eval() + self.model.half() # Enable FP16 for better performance + torch.set_float32_matmul_precision('high') + self.model.to(device) + self.current_model_version = model_name + + except Exception as e: + handle_model_error(f"Error loading BiRefNet model: {str(e)}") + + def process_image(self, images, model_name, params): + try: + self.load_model(model_name) + + if isinstance(images, torch.Tensor): + if len(images.shape) == 3: + images = [images] + else: + images = [img for img in images] + + all_masks = [] + + transform_image = transforms.Compose([ + transforms.Resize((2048, 2048)), + transforms.ToTensor(), + transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) + ]) + + for img in images: + orig_image = tensor2pil(img) + w, h = orig_image.size + + input_tensor = transform_image(orig_image).unsqueeze(0).to(device).half() + + with torch.no_grad(): + pred = self.model(input_tensor)[-1].sigmoid().cpu() + + pred = pred[0].squeeze() + mask = transforms.ToPILImage()(pred) + mask = mask.resize((w, h), Image.LANCZOS) + + all_masks.append(mask) + torch.cuda.empty_cache() + + return all_masks[0] if len(all_masks) == 1 else all_masks + + except Exception as e: + handle_model_error(f"Error in BiRefNet processing: {str(e)}") + class RMBG: def __init__(self): self.models = { "RMBG-2.0": RMBGModel(), "INSPYRENET": InspyrenetModel(), "BEN": BENModel(), - "BEN2": BEN2Model() + "BEN2": BEN2Model(), + "BIREFNET-HR": BiRefNetModel() } @classmethod