Add files via upload
This commit is contained in:
AI Lab
+106
-5
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user